用pytorch写一段代码实现mnist手写数字识别

import torch
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms# 载入MNIST数据集
train_dataset = dsets.MNIST(root='./data', train=True,
transform=transforms.ToTensor(), download=True)
test_dataset = dsets.MNIST(root='./data', train=False,
transform=transforms.ToTensor())# 创建DataLoader加载数据
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=100, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=100, shuffle=False)# 创建模型
class CNNModel(nn.Module):
def __init__(self):
super(CNNModel, self).__init__()
# 卷积层
self.cnn1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2)
self.relu1 = nn.ReLU()
# 最大池化层
self.maxpool1 = nn.MaxPool2d(kernel_size=2)
# 全连接层
self.fc1 = nn.Linear(16*7*7, 10)

def forward(self, x):
# 卷积、激活、最大池化
out = self.cnn1(x)
out = self.relu1(out)
out = self.maxpool1(out)
# 扁平化
out = out.view(out.size(0), -1)
# 全连接层
out = self.fc1(out)
return out# 实例化模型
model = CNNModel()# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)# 训练模型
for epoch in range(2):
for i, (images, labels) in enumerate(train_loader):
# 前向传播
outputs = model(images)
# 计算损失函数
loss = criterion(outputs, labels)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()

if (i+1) % 100 == 0:
print ('Epoch : %d/%d, Iter : %d/%d, Loss: %.4f'
%(epoch+1, 2, i+1, len(train_dataset)//100, loss.item()))# 测试模型
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_loader:
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum()print('10000张图片的准确率为: %d %%' % (100 * correct / total))