代码是根据《深度学习之Pytorch实战计算机视觉》第六章实战手写体识别,经自己的测试思考改编而来。
文章源码可参考:
https://github.com/JaimeTang/book-code/blob/master/chapter-6/chapter-6.ipynb
普通的cpu训练迭代一轮需要5分钟左右,(一共5轮)测试了几次觉得过于浪费时间。
更换GPU训练5轮一共需要3-5分钟,非常方便。
代码一共修改了四个地方,见代码块。
#MNIST实战 GPU版
#具体来说需要更改的地方有四个:
import torch
import torchvision
from torchvision import datasets,transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
transform=transforms.Compose([transforms.ToTensor(),
transforms.Normalize(mean=[0.5],std=[0.5])])
data_train = datasets.MNIST(root = "./data/", transform = transform, train = True, download = True)
data_test = datasets.MNIST(root = "./data/", transform = transform, train = False)
dat