基于GPU的Pytorch——CNN实现MNIST数据识别

CNN实现MNIST数据集分类时，发现CPU训练特别慢，20个epoch需要半个多小时，故修改为GPU训练，训练时长明显缩短。

修改如下：

1. 需要把网络结构移到GPU上：

   model = Net()
   if torch.cuda.is_available():
       print("The model is using GPU")
       model = model.cuda()
   

2. 训练数据集和测试数据集移到GPU上：

   for i, data in enumerate(train_loader):
           # 获得数据和对应的标签
           inputs, labels = data
           # 获得模型预测结果，（64，10）
           inputs = inputs.cuda()
           labels = labels.cuda()
           out = model(inputs)
           # 交叉熵代价函数out(batch,C),labels(batch)
           loss = entropy_loss(out, labels)
   

   注意： 刚开始由于只修改了数据，未修改标签，导致报错：
   RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 ‘mat1’ in call to _th_addmm

完整代码如下——使用工具：jupyter notebook：

#%%
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import  matplotlib.pyplot as plt

#%%

# 下载训练集
train_dataset = datasets.MNIST(root='./',
                train=True,
                transform=transforms.ToTensor(),
                download=False)
# 下载测试集
test_dataset = datasets.MNIST(root='./',
               train=False,
               transform=transforms.ToTensor(),
               download