【Pytorch】MNIST数据集上卷积神经网络的实现

1.配置库和配置参数

import torch
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms 
from torch.autograd import Variable

# 配置参数
torch.manual_seed(1) #设置随机数种子，确保结果可重复
input_size = 784
hidden_size = 500
num_classes = 10
num_epoched = 5
batch_size = 100
learning_rate = 0.001

2.加载MNIST数据

train_dataset = dsets.MNIST(
	root = './data',  # 数据保持的位置
	train = True,  #训练集
	transform = transforms.ToTensor(),   #一个取值范围[0,255]的PIL.Image，转化为[0, 1.0]的torch.FloatTensor
	download = True     # 下载数据
)

test_dataset = dsets.MNIST(
	root = './data',
	train = True,  #测试集
	transform = transforms.ToTensor()
)

3.数据的批处理

train_loader = torch.utils.data.DataLoader(
	datasets = train_dataset,
	batch_size = batch_size,
	shuffle = True
	)

test_loader = torch.utils.data.DataLoader(
	datasets = test_dataset,
	batch_size = batch_size,
	shuffle = False
	)

4.创建CNN模型

Class CNN(nn.Module):
	def __init__(self, in_dim, n_classes): # 28*28*1
		super(Cnn, self).__init__()
		self.conv = nn.Sequential(
			nn.Conv2d(in_dim, 6, 3, stride=1, padding=1),   #  in_channels, out_channels, kernel_size, stride, padding  28*28*6
			nn.RelU(True),
			nn.MaxPool2d(2,2),   # kernel_size, stride 14*14*6
			nn.Conv2d(6, 16, 5, stride=1, padding=0)   # 10*10*16
			nn.ReLU(True),
			nn.MaxPooling2d(2,2) ) #5*5*16

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

	def forward(self, x):
		out = self.conv(x)
		out = out.view(out.size(0), 400)
		out = self.fc(out)
		return out

net = Cnn(1, 10)

5.模型训练

# 定义loss和optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr = learning rate)

for epoch in range(num_epoches):
	running_loss = 0.0
	running_acc = 0.0
	for i, data in emurate(train_loader, 1):
		img, lable = data
		img, label = Variable(img), Variable(label)

		#	前向传播
		out = net(img)
		loss = criterion(out, label)
		running_loss += loss.data[0] * label.size(0) # total_loss, 由于loss是batch取均值的，需要把batch_size乘回去
		_, pred = torch.max(out, 1) # softmax后，在第1个维度上取最大值
		num_correct = (pred==label).sum()
		running_acc += num_correct.data[0]  #正确结果的总数

		# 后向传播
		optimizer.zero_grad()
		loss.backward()
		optimizer.step() 

6.在测试集上测试识别率

net.eval() #由于训练和测试BatchNorm，Dropout配置不同，需要说明是否模型测试
correct = 0
total = 0
for images, lables in test_loader:
	images =  = Variable(image.view(-1, 28*28))
	
	outputs = net(images)
	_, predicted = torch.max(outputs.data, 1)  #按维度1返回最大值
	total += labels.size(0)   # 正确结果
	correct += (predicted == labels).sum()   # 正确结果总数

参考资料：PyTorch机器学习从入门到实战