手写数字识别

复制代码后依次运行下面命令

python3 lenet5.py python3 train.py python3 test.py 

lenet5.py

import torch.nn as nn  class LeNet5(nn.Module): 	def __init__(self): 		super(LeNet5, self).__init__() 		self.conv1 = nn.Sequential(	 			nn.Conv2d(				# (1, 28, 28) 				in_channels=1, 				out_channels=6, 				kernel_size=5, 				stride=1, 				padding=2 			),						# (6, 28, 28) 			nn.ReLU(), 			nn.MaxPool2d(kernel_size=2, stride=2),	# (6, 14, 14)	 		) 		self.conv2 = nn.Sequential( 			nn.Conv2d(6, 16, 5, 1, 0), 			nn.ReLU(),				# (16, 10, 10) 			nn.MaxPool2d(2, 2)		# (16, 5, 5) 		) 		self.LinNet = nn.Sequential( 			nn.Linear(16 * 5 * 5, 120), 			nn.ReLU(), 			nn.Linear(120, 84), 			nn.ReLU(), 			nn.Linear(84, 10) 		)  	def forward(self, x): 		x = self.conv1(x) 		x = self.conv2(x) 		x = x.view(x.size(0), -1) 		out = self.LinNet(x) 		return out  if __name__ == __main__: 	myNet = LeNet5() 	print(myNet) 

train.py

from lenet5 import LeNet5 import time import torch import torch.nn as nn import numpy as np import matplotlib.pyplot as plt from torch.utils.data import DataLoader from torchvision import transforms from torchvision.datasets import MNIST  transform = transforms.Compose([ 	transforms.ToTensor(), 	transforms.Normalize((0.1307, ), (0.3081, )) ])  train_set = MNIST(root='./MNIST',                   train=True,                   download=True,                   transform=transform)  train_loader = DataLoader(     train_set,     batch_size = 100,     shuffle = True )  net = LeNet5().cuda() optimizer = torch.optim.SGD(net.parameters(), lr=0.01) loss_func = nn.CrossEntropyLoss()  start_time = time.time()  epochs = 10 epoch_loss = [] for epoch in range(epochs): 	running_loss = 0.0 	for i, (inputs, labels) in enumerate(train_loader): 		inputs = torch.tensor(inputs).type(torch.FloatTensor).cuda() 		labels = torch.tensor(labels).type(torch.LongTensor).cuda()  		out = net(inputs) 		loss = loss_func(out, labels) 		optimizer.zero_grad() 		loss.backward() 		optimizer.step()  		running_loss += loss.item() 		avr_loss = running_loss / (i+1) 	epoch_loss.append(avr_loss) 	print('epoch %d, loss: %.3f' % (epoch, avr_loss))  end_time = time.time() print('Finished training, time used: %.3f' % (end_time - start_time))  plt.figure(figsize=(8, 5), dpi=150) plt.plot(epoch_loss, c='r') plt.savefig('./document/figure/loss.pdf') plt.show() torch.save(net, 'net.pkl') 

test.py

import torch from torch.utils.data import DataLoader from torchvision import transforms from torchvision.datasets import MNIST  transform = transforms.Compose([ 	transforms.ToTensor(), 	transforms.Normalize((0.1307, ), (0.3081, )) ])  test_set = MNIST(root='./MNIST',                 train=False,                 download=True,                 transform=transform)  test_loader = DataLoader(     test_set,     batch_size = 100,     shuffle = True )  net = torch.load('net.pkl')  correct = 0 for i, (inputs, labels) in enumerate(test_loader): 	inputs = torch.tensor(inputs).type(torch.FloatTensor).cuda() 	labels = torch.tensor(labels).type(torch.LongTensor).cuda()  	out = net(inputs) 	y_pred = out.argmax(axis=1)  	correct += (y_pred == labels).sum()  	print('Average accuracy of %d tests: %.1f' % ((i+1), correct / (i+1)))