Pytorch-实战MNIST-GPU版

最新推荐文章于 2024-09-08 15:42:39 发布

土耳其的曼谷风

最新推荐文章于 2024-09-08 15:42:39 发布

阅读量785

点赞数 2

本文链接：https://blog.csdn.net/ihiefoxboq/article/details/107295202

版权

本文基于《深度学习之Pytorch实战计算机视觉》的第六章，介绍了如何使用GPU进行MNIST手写体识别的实战。原始CPU训练耗时较长，而经过GPU优化后，训练速度显著提升，5轮训练只需3-5分钟。文中提到了四个关键的代码修改部分。

摘要由CSDN通过智能技术生成

代码是根据《深度学习之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