Task7.手写数字识别

用PyTorch完成手写数字识别

 import numpy as np
 import torch
 from torch import nn, optim
 import torch.nn.functional as F
 from torch.autograd import Variable
 from torch.utils.data import DataLoader
 from torchvision import transforms
 from torchvision import datasets

 batch_size = 128
 learning_rate = 0.01
 num_epoch = 10

 # 实例化MNIST数据集对象
 train_data = datasets.MNIST('./dataset', train=True, transform=transforms.ToTensor(), download=True)
 test_data = datasets.MNIST('./dataset', train=False, transform=transforms.ToTensor(), download=True)

 # train_loader：以batch_size大小的样本组为单位的可迭代对象
 train_loader = DataLoader(train_data, batch_size, shuffle=True)
 test_loader = DataLoader(test_data)

 class CNN(nn.Module):
     def __init__(self, in_dim, out_dim):
         super(CNN, self).__init__()
         self.conv1 = nn.Conv2d(in_dim, 6, 3, stride=1, padding=1)
         self.batch_norm1 = nn.BatchNorm2d(6)
         self.relu = nn.ReLU(True)
         self.conv2 = nn.Conv2d(6, 16, 5, stride=1, padding=0)
         self.pool = nn.MaxPool2d(2, 2)
         self.batch_norm2 = nn.BatchNorm2d(16)

         self.fc1 = nn.Linear(400, 120)
         self.fc2 = nn.Linear(120, 84)
         self.fc3 = nn.Linear(84, out_dim)

     def forward(self, x):
         x = self.batch_norm1(self.conv1(x))
         x = F.relu(x)
         x = self.pool(x)
         x = self.batch_norm2(self.conv2(x))
         x = self.relu(x)
         x = self.pool(x)
         x = x.view(x.size(0), -1)
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))
         x = self.fc3(x)
         return x

     def print_model_name(self):
         print("Model Name: CNN")

 class Cnn(nn.Module):
     def __init__(self, in_dim, n_class):
         super(Cnn, self).__init__()
         self.conv = nn.Sequential(
             nn.Conv2d(in_dim, 6, 3, stride=1, padding=1),
             nn.ReLU(True),
             nn.MaxPool2d(2, 2),
             nn.Conv2d(6, 16, 5, stride=1, padding=0),
             nn.ReLU(True),
             nn.MaxPool2d(2, 2))

         self.fc = nn.Sequential(
             nn.Linear(400, 120), nn.Linear(120, 84), nn.Linear(84, n_class))

     def forward(self, x):
         # print(x.size()) torch.Size([1024, 1, 28, 28])
         out = self.conv(x)
         out = out.view(out.size(0), -1)
         # print(out.size()) = torch.Size([1024, 400])
         out = self.fc(out)
         # print(out.size()) torch.Size([1024, 10])
         return out

     def print_model_name(self):
         print("Model Name: Cnn")

 isGPU = torch.cuda.is_available()
 print(isGPU)
 model = CNN(1, 10)
 if isGPU:
     model = model.cuda()
 criterion = nn.CrossEntropyLoss()
 optimizer = optim.SGD(model.parameters(), lr=learning_rate)
 for epoch in range(num_epoch):
     running_acc = 0.0
     running_loss = 0.0
     for i, data in enumerate(train_loader, 1): # train_loader：以batch_size大小的样本组为单位的可迭代对象
         img, label = data
         img = Variable(img)
         label = Variable(label)
         if isGPU:
             img = img.cuda()
             label = label.cuda()
         # forward
         out = model(img)
         loss = criterion(out, label)
         # print(label)
         # backward
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()

         _, pred = torch.max(out, dim=1)  # 按维度dim 返回最大值
         running_loss += loss.item()*label.size(0)
         current_num = (pred == label).sum() # variable
         acc = (pred == label).float().mean()        # variable
         running_acc += current_num.item()

         if i % 100 == 0:
             print("epoch: {}/{}, loss: {:.6f}, running_acc: {:.6f}"
                   .format(epoch+1, num_epoch, loss.item(), acc.item()))
     print("epoch: {}, loss: {:.6f}, accuracy: {:.6f}".format(epoch+1, running_loss, running_acc/len(train_data)))

 model.eval()
 current_num = 0
 for i , data in enumerate(test_loader, 1):
     img, label = data
     if isGPU:
         img = img.cuda()
         label = label.cuda()
     with torch.no_grad():
         img = Variable(img)
         label = Variable(label)
     out = model(img)
     _, pred = torch.max(out, 1)
     current_num += (pred == label).sum().item()

 print("Test result: accuracy: {:.6f}".format(float(current_num/len(test_data))))

 torch.save(model.state_dict(), './cnn.pth') # 保存模型