代码如下
# -*- coding: utf-8 -*-
import torch
from torch import nn
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.optim as optim
#定义神经网络类Network, 它将继承nn.Module类
class Network(nn.Module):
# 神经网络中的神经元数量是固定的,所以init不用传入参数
def __init__(self):
super().__init__() # 调用了父类的初始化函数
# layer1是输入层于隐藏层之间的线性层
self.layer1 = nn.Linear(28 * 28, 256)
# layer2是隐藏层和输出层之间的线性层
self.layer2 = nn.Linear(256, 10)
def forward(self, x):
#将输入由二维转换为一维,便于输入到线性层中
x = x.view(-1, 28*28)
x = self.layer1(x) #输出layer1的输出结果
x = torch.relu(x) #输出激活后的输出结果
#计算layer2的结果,并返回
return self.layer2(x)
device = torch.device("cuda:0")
# 加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())
# 定义数据加载器
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=128, shuffle=False)
# 定义模型、损失函数和优化器
model = Network().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
# 训练模型
for epoch in range(20):
for i, (images, labels) in enumerate(train_loader):
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if (i+1) % 100 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, 20, i+1, len(train_loader), loss.item()))
# 测试模型
model.eval()
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Test Accuracy: {:.2f}%'.format(100 * correct / total))
测试集上的准确率还是蛮高的,达到了98.16%。