Pytorch-实战MNIST-GPU版

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

data_loader_train = torch.utils.data.DataLoader(dataset = data_train,batch_size = 64, shuffle = True)
data_loader_test = torch.utils.data.DataLoader(dataset = data_test, batch_size = 64,shuffle = True)

images , labels = next(iter(data_loader_train))
img = torchvision.utils.make_grid(images)

img = img.numpy().transpose(1,2,0)
std = [0.5]
mean = [0.5]
img = img * std +mean
print([labels[i] for i in range(64)])
plt.imshow(img)

class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1=torch.nn.Sequential(
            torch.nn.Conv2d(1,64,kernel_size=3,stride=1,padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(stride=2,kernel_size=2))
        self.dense = torch.nn.Sequential(torch.nn.Linear(14*14*128,1024), torch.nn.ReLU(),torch.nn.Dropout(p = 0.5),torch.nn.Linear(1024,10))
        
    def forward(self, x):
        x = self.conv1(x)
        x = x.view(-1, 14*14*128)
        x = self.dense(x)
        return x
 
model = Model()
model = model.cuda()                                       #第一处修改
print("Model= ",model)
cost =torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

n_epochs = 5

for epoch in range(n_epochs):
    running_loss = 0.0
    running_correct = 0
    print("Epoch {}/{}".format(epoch,n_epochs))
    print("-"*10)
    for data in data_loader_train:
        x_train, y_train = data
        x_train, y_train = Variable(x_train.cuda()),Variable(y_train.cuda()) 
        #第二处修改 cpu版->  Variable(x_train),Variable(y_train) 

        outputs = model(x_train)
        _,pred = torch.max(outputs.data,1)
        optimizer.zero_grad()
        loss = cost(outputs, y_train)
        
        loss.backward()
        optimizer.step()
        running_loss += loss.data
        running_correct += torch.sum(pred == y_train.data)

    testing_correct = 0
    for data in data_loader_test:
        x_test,y_test = data
        x_test,y_test = Variable(x_test.cuda()),Variable(y_test.cuda())
        #第三处更改 加上.cuda() 同上

        outputs = model(x_test)
        _,pred = torch.max(outputs.data,1)
        testing_correct += torch.sum(pred == y_test.data)

    print("Loss is :{:.4f},Train Accuracy is : {:.4f}%, Test Accuracy is {:.4f}".format(running_loss/len(data_train),100*running_correct/len(data_train),100*testing_correct/len(data_test)))
    
data_loader_test = torch.utils.data.DataLoader(dataset = data_test,batch_size = 4, shuffle = True)
x_test , y_test = next(iter(data_loader_test))
inputs = Variable(x_test.cuda())  
#第四处更改 加上.cuda() 此位置不改最后测试输出type对应不上

pred = model(inputs)
_,pred = torch.max(pred,1)
print("Predict Label is:",[ i for i in pred.data])
print("Real Label is:",[i for i in y_test])
img = torchvision.utils.make_grid(x_test)
img = img.numpy().transpose(1,2,0)
std = [0.5,0.5,0.5]
mean = [0.5,0.5,0.5]
img = img*std + mean
plt.imshow(img)