用PyTorch完成手写数字识别
1 import numpy as np 2 import torch 3 from torch import nn, optim 4 import torch.nn.functional as F 5 from torch.autograd import Variable 6 from torch.utils.data import DataLoader 7 from torchvision import transforms 8 from torchvision import datasets 9 10 batch_size = 128 11 learning_rate = 0.01 12 num_epoch = 10 13 14 # 实例化MNIST数据集对象 15 train_data = datasets.MNIST('./dataset', train=True, transform=transforms.ToTensor(), download=True) 16 test_data = datasets.MNIST('./dataset', train=False, transform=transforms.ToTensor(), download=True) 17 18 # train_loader：以batch_size大小的样本组为单位的可迭代对象 19 train_loader = DataLoader(train_data, batch_size, shuffle=True) 20 test_loader = DataLoader(test_data) 21 22 class CNN(nn.Module): 23 def __init__(self, in_dim, out_dim): 24 super(CNN, self).__init__() 25 self.conv1 = nn.Conv2d(in_dim, 6, 3, stride=1, padding=1) 26 self.batch_norm1 = nn.BatchNorm2d(6) 27 self.relu = nn.ReLU(True) 28 self.conv2 = nn.Conv2d(6, 16, 5, stride=1, padding=0) 29 self.pool = nn.MaxPool2d(2, 2) 30 self.batch_norm2 = nn.BatchNorm2d(16) 31 32 self.fc1 = nn.Linear(400, 120) 33 self.fc2 = nn.Linear(120, 84) 34 self.fc3 = nn.Linear(84, out_dim) 35 36 def forward(self, x): 37 x = self.batch_norm1(self.conv1(x)) 38 x = F.relu(x) 39 x = self.pool(x) 40 x = self.batch_norm2(self.conv2(x)) 41 x = self.relu(x) 42 x = self.pool(x) 43 x = x.view(x.size(0), -1) 44 x = F.relu(self.fc1(x)) 45 x = F.relu(self.fc2(x)) 46 x = self.fc3(x) 47 return x 48 49 def print_model_name(self): 50 print("Model Name: CNN") 51 52 53 class Cnn(nn.Module): 54 def __init__(self, in_dim, n_class): 55 super(Cnn, self).__init__() 56 self.conv = nn.Sequential( 57 nn.Conv2d(in_dim, 6, 3, stride=1, padding=1), 58 nn.ReLU(True), 59 nn.MaxPool2d(2, 2), 60 nn.Conv2d(6, 16, 5, stride=1, padding=0), 61 nn.ReLU(True), 62 nn.MaxPool2d(2, 2)) 63 64 self.fc = nn.Sequential( 65 nn.Linear(400, 120), nn.Linear(120, 84), nn.Linear(84, n_class)) 66 67 def forward(self, x): 68 # print(x.size()) torch.Size([1024, 1, 28, 28]) 69 out = self.conv(x) 70 out = out.view(out.size(0), -1) 71 # print(out.size()) = torch.Size([1024, 400]) 72 out = self.fc(out) 73 # print(out.size()) torch.Size([1024, 10]) 74 return out 75 76 def print_model_name(self): 77 print("Model Name: Cnn") 78 79 80 isGPU = torch.cuda.is_available() 81 print(isGPU) 82 model = CNN(1, 10) 83 if isGPU: 84 model = model.cuda() 85 criterion = nn.CrossEntropyLoss() 86 optimizer = optim.SGD(model.parameters(), lr=learning_rate) 87 for epoch in range(num_epoch): 88 running_acc = 0.0 89 running_loss = 0.0 90 for i, data in enumerate(train_loader, 1): # train_loader：以batch_size大小的样本组为单位的可迭代对象 91 img, label = data 92 img = Variable(img) 93 label = Variable(label) 94 if isGPU: 95 img = img.cuda() 96 label = label.cuda() 97 # forward 98 out = model(img) 99 loss = criterion(out, label) 100 # print(label) 101 # backward 102 optimizer.zero_grad() 103 loss.backward() 104 optimizer.step() 105 106 _, pred = torch.max(out, dim=1) # 按维度dim 返回最大值 107 running_loss += loss.item()*label.size(0) 108 current_num = (pred == label).sum() # variable 109 acc = (pred == label).float().mean() # variable 110 running_acc += current_num.item() 111 112 if i % 100 == 0: 113 print("epoch: {}/{}, loss: {:.6f}, running_acc: {:.6f}" 114 .format(epoch+1, num_epoch, loss.item(), acc.item())) 115 print("epoch: {}, loss: {:.6f}, accuracy: {:.6f}".format(epoch+1, running_loss, running_acc/len(train_data))) 116 117 model.eval() 118 current_num = 0 119 for i , data in enumerate(test_loader, 1): 120 img, label = data 121 if isGPU: 122 img = img.cuda() 123 label = label.cuda() 124 with torch.no_grad(): 125 img = Variable(img) 126 label = Variable(label) 127 out = model(img) 128 _, pred = torch.max(out, 1) 129 current_num += (pred == label).sum().item() 130 131 print("Test result: accuracy: {:.6f}".format(float(current_num/len(test_data)))) 132 133 torch.save(model.state_dict(), './cnn.pth') # 保存模型