首先,我们导入所需的库:
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import matplotlib.pyplot as plt
然后,我们加载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())
接下来,我们定义MLP模型类:
class MLP(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(MLP, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out
在这个例子中,MLP模型有两个全连接层。在初始化函数__init__()
中,我们定义了两个全连接层self.fc1
和self.fc2
,激活函数ReLUself.relu
。在前向传播函数forward()
中,我们将输入数据x
传递给fc1
,然后通过ReLU激活函数,再传递给fc2
。
我们继续设置超参数:
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001
这里的input_size
是输入图像的大小(28x28=784),hidden_size
是隐藏层的大小,num_classes
是输出的类别数(这里是10个数字),num_epochs
是训练的轮数,batch_size
是每批训练样本的数量,learning_rate
是学习率。
然后,我们准备数据加载器:
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
数据加载器用于将数据分成小批次进行训练和测试。
接下来,我们实例化模型、定义损失函数和优化器:
model = MLP(input_size, hidden_size, num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
在每个训练周期中,我们对训练数据进行迭代,进行前向传播、计算损失、反向传播和参数更新:
total_step = len(train_loader)
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.reshape(-1, 28*28)
# 前向传播和计算损失
outputs = model(images)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i+1) % 100 == 0:
print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{total_step}], Loss: {loss.item()}')
在每个训练周期结束后,我们在测试数据上进行模型评估,计算准确率:
model.eval()
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
images = images.reshape(-1, 28*28)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'该模型在10000张测试图像上的准确性: {100 * correct / total}%')
最后,我们显示一些图像示例:
images, labels = next(iter(test_loader))
images = images[:5]
labels = labels[:5]
fig, axes = plt.subplots(1, 5, figsize=(10,2))
for i in range(5):
image = images[i].numpy().squeeze()
axes[i].imshow(image, cmap='gray')
axes[i].axis('off')
axes[i].set_title(f'Ground Truth: {labels[i]}')
plt.show()
运行结果展示:
Epoch [1/5], Step [100/600], Loss: 0.5327053666114807
Epoch [1/5], Step [200/600], Loss: 0.1780116856098175
Epoch [1/5], Step [300/600], Loss: 0.22097641229629517
Epoch [1/5], Step [400/600], Loss: 0.21554379165172577
Epoch [1/5], Step [500/600], Loss: 0.28248608112335205
Epoch [1/5], Step [600/600], Loss: 0.08389710634946823
Epoch [2/5], Step [100/600], Loss: 0.15392127633094788
Epoch [2/5], Step [200/600], Loss: 0.139845073223114
Epoch [2/5], Step [300/600], Loss: 0.10855261981487274
Epoch [2/5], Step [400/600], Loss: 0.05962827056646347
Epoch [2/5], Step [500/600], Loss: 0.0902574360370636
Epoch [2/5], Step [600/600], Loss: 0.153961181640625
Epoch [3/5], Step [100/600], Loss: 0.08641496300697327
Epoch [3/5], Step [200/600], Loss: 0.05064410716295242
Epoch [3/5], Step [300/600], Loss: 0.05174357816576958
Epoch [3/5], Step [400/600], Loss: 0.05410122871398926
Epoch [3/5], Step [500/600], Loss: 0.07135355472564697
Epoch [3/5], Step [600/600], Loss: 0.05457733944058418
Epoch [4/5], Step [100/600], Loss: 0.07079239934682846
Epoch [4/5], Step [200/600], Loss: 0.048895031213760376
Epoch [4/5], Step [300/600], Loss: 0.09586360305547714
Epoch [4/5], Step [400/600], Loss: 0.03884414583444595
Epoch [4/5], Step [500/600], Loss: 0.127535879611969
Epoch [4/5], Step [600/600], Loss: 0.024616247043013573
Epoch [5/5], Step [100/600], Loss: 0.024639535695314407
Epoch [5/5], Step [200/600], Loss: 0.006189913023263216
Epoch [5/5], Step [300/600], Loss: 0.08351482450962067
Epoch [5/5], Step [400/600], Loss: 0.024257291108369827
Epoch [5/5], Step [500/600], Loss: 0.032762959599494934
Epoch [5/5], Step [600/600], Loss: 0.016724038869142532
该模型在10000张测试图像上的准确性: 97.86%
整体代码:
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import matplotlib.pyplot as plt
# 加载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())
# 定义模型
class MLP(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(MLP, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out
# 设置超参数
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001
# 准备数据加载器
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
# 定义模型和损失函数
model = MLP(input_size, hidden_size, num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# 训练模型
total_step = len(train_loader)
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.reshape(-1, 28 * 28)
# 前向传播和计算损失
outputs = model(images)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0:
print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{total_step}], Loss: {loss.item()}')
# 在测试集上测试模型
model.eval()
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
images = images.reshape(-1, 28 * 28)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'该模型在10000张测试图像上的准确性: {100 * correct / total}%')
# 显示图像示例
images, labels = next(iter(test_loader))
images = images[:5]
labels = labels[:5]
fig, axes = plt.subplots(1, 5, figsize=(10, 2))
for i in range(5):
image = images[i].numpy().squeeze()
axes[i].imshow(image, cmap='gray')
axes[i].axis('off')
axes[i].set_title(f'Ground Truth: {labels[i]}')
plt.show()