記得第一次接觸手寫數字識別數據集還在學習TensorFlow，各種sess.run()，頭都繞暈了。自從接觸pytorch以來，一直想寫點什么。曾經在2017年5月，Andrej Karpathy發表的一篇Twitter，調侃道：l've been using PyTorch a few months now, l've never felt better, l've more energy.My skin is clearer. My eye sight has improved。確實，使用pytorch以來，確實感覺心情要好多了，不像TensorFlow那樣晦澀難懂。迫不及待的用pytorch實戰了一把MNIST數據集，構建LeNet神經網絡。話不多說，直接上代碼！
import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets,transforms import torchvision from torch.autograd import Variable from torch.utils.data import DataLoader import cv2 class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(1, 6, 3, 1, 2), nn.ReLU(), nn.MaxPool2d(2, 2) ) self.conv2 = nn.Sequential( nn.Conv2d(6, 16, 5), nn.ReLU(), nn.MaxPool2d(2, 2) ) self.fc1 = nn.Sequential( nn.Linear(16 * 5 * 5, 120), nn.BatchNorm1d(120), nn.ReLU() ) self.fc2 = nn.Sequential( nn.Linear(120, 84), nn.BatchNorm1d(84),#加快收斂速度的方法（注：批標准化一般放在全連接層后面，激活函數層的前面） nn.ReLU() ) self.fc3 = nn.Linear(84, 10) # self.sfx = nn.Softmax() def forward(self, x): x = self.conv1(x) x = self.conv2(x) # print(x.shape) x = x.view(x.size()[0], -1) x = self.fc1(x) x = self.fc2(x) x = self.fc3(x) # x = self.sfx(x) return x device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') batch_size = 64 LR = 0.001 Momentum = 0.9 # 下載數據集 train_dataset = datasets.MNIST(root = './data/', train=True, transform = transforms.ToTensor(), download=False) test_dataset =datasets.MNIST(root = './data/', train=False, transform=transforms.ToTensor(), download=False) #建立一個數據迭代器 train_loader = torch.utils.data.DataLoader(dataset = train_dataset, batch_size = batch_size, shuffle = True) test_loader = torch.utils.data.DataLoader(dataset = test_dataset, batch_size = batch_size, shuffle = False) #實現單張圖片可視化 # images,labels = next(iter(train_loader)) # img = torchvision.utils.make_grid(images) # img = img.numpy().transpose(1,2,0) # # img.shape # std = [0.5,0.5,0.5] # mean = [0.5,0.5,0.5] # img = img*std +mean # cv2.imshow('win',img) # key_pressed = cv2.waitKey(0) net = LeNet().to(device) criterion = nn.CrossEntropyLoss()#定義損失函數 optimizer = optim.SGD(net.parameters(),lr=LR,momentum=Momentum) epoch = 1 if __name__ == '__main__': for epoch in range(epoch): sum_loss = 0.0 for i, data in enumerate(train_loader): inputs, labels = data inputs, labels = Variable(inputs).cuda(), Variable(labels).cuda() optimizer.zero_grad()#將梯度歸零 outputs = net(inputs)#將數據傳入網絡進行前向運算 loss = criterion(outputs, labels)#得到損失函數 loss.backward()#反向傳播 optimizer.step()#通過梯度做一步參數更新 # print(loss) sum_loss += loss.item() if i % 100 == 99: print('[%d,%d] loss:%.03f' % (epoch + 1, i + 1, sum_loss / 100)) sum_loss = 0.0 #驗證測試集 net.eval()#將模型變換為測試模式 correct = 0 total = 0 for data_test in test_loader: images, labels = data_test images, labels = Variable(images).cuda(), Variable(labels).cuda() output_test = net(images) # print("output_test:",output_test.shape) _, predicted = torch.max(output_test, 1)#此處的predicted獲取的是最大值的下標 # print("predicted:",predicted.shape) total += labels.size(0) correct += (predicted == labels).sum() print("correct1: ",correct) print("Test acc: {0}".format(correct.item() / len(test_dataset)))#.cpu().numpy()
本次識別手寫數字，只做了1個epoch，train_loss：0.250，測試集上的准確率：0.9685，相當不錯的結果。
