Pytorch搭建MyNet实现MNIST手写数字识别
视频：https://www.bilibili.com/video/BV1Wf421B74f/?spm_id_from=333.880.my_history.page.click
1.1 Model类
import torch
import torch.nn as nn
# 改进的三层神经网络
class MyNet(nn.Module):
def __init__(self):
super().__init__()
# 定义全连接层
self.fc1 = nn.Linear(28 * 28, 256) # 输入层 输入是28*28的灰度图像，输出是256个神经元
self.fc2 = nn.Linear(256, 128) # 第二层，全连接层，输入256个神经元，输出128个神经元
self.fc3 = nn.Linear(128, 64) # 第三层，全连接层，输入128个神经元，输出64个神经元
self.fc4 = nn.Linear(64, 10) # 第四层，全连接层，输入64个神经元，输出10个类别
def forward(self, x):
x = torch.flatten(x, start_dim=1) # 展平数据，方便进行全连接
x = torch.relu(self.fc1(x)) # 第一层+ReLU激活
x = torch.relu(self.fc2(x)) # 第二层+ReLU激活
x = torch.relu(self.fc3(x)) # 第三层+ReLU激活
x = self.fc4(x) # 第四层 输出层 不需要激活函数
return x
2 train.py创建
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from .model import MyNet
# 设置随机种子
torch.manual_seed(21)
# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device:{device}")
transform = transforms.Compose([
transforms.ToTensor(), # 将图像转为张量
transforms.Normalize((0.5,), (0.5,)) # 标准化图像数据，对于灰度图像，只需要一个通道的标准化
])
# 加载MNIST数据集
train_dataset = datasets.MNIST(root='../../datasets', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='../../datasets', train=False, download=True, transform=transform)
# 创建数据加载器
batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 初始化模型并将模型移到GPU上
model = MyNet().to(device)
# 定义损失函数和优化器
lr = 0.001 # 学习率为0.001
criterion = nn.CrossEntropyLoss() # 交叉熵
optimizer = optim.Adam(model.parameters(), lr=lr) # 使用Adam优化器
# 保存训练过程中的损失和准确率
train_losses = []
train_accuracies = []
test_accuracies = []
epochs = 10
best_accuracy = 0.0 # 记录最佳验证集准确率
best_model_path = 'best_mnist_model.pth'
# 训练10个epoch
for epoch in range(epochs):
running_loss = 0.0
corrct_train = 0 # 正确预测的数量
total_train = 0 # 样本总数
# 训练过程
model.train() # 设定模型为训练模式
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device) # 将数据移动到GPU上
optimizer.zero_grad() # 梯度清零
outputs = model(inputs) # 前向传播
loss = criterion(outputs, labels) # 计算损失
loss.backward() # 反向传播
optimizer.step() # 更新参数
running_loss += loss.item() # 累加损失
# 计算训练集上的准确率
_, predicted = torch.max(outputs, 1) # 获取预测结果
total_train += labels.size(0) # 累加样本数量
corrct_train += (predicted == labels).sum().item() # 累加正确预测的数量
# 计算训练集上的准确率
train_accuracy = corrct_train / total_train
train_losses.append(running_loss / len(train_loader)) # 记录每个 epoch 的平均损失,len(train_loader)为批次数，一个epoch结束后
train_accuracies.append(train_accuracy) # 记录每个 epoch 的训练集准确率
print(
f"Epoch: {epoch + 1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}, Train Accuracy: {train_accuracy:.2%}")
# 在测试集上评估模型
model.eval() # 设定模型为评估模式
correct = 0 # 正确的预测数量
total = 0 # 样本总数
with torch.no_grad(): # 关闭梯度计算
for inputs, labels in test_loader:
inputs, labels = inputs.to(device), labels.to(device) # 将数据移动到GPU上
outputs = model(inputs) # 前向传播
_, predicted = torch.max(outputs, 1) # 获取预测结果
total += labels.size(0) # 累加样本数量
correct += (predicted == labels).sum().item() # 累加正确预测的数量
# 计算测试集上的准确率
test_accuracy = correct / total # 记录每个 epoch 的测试集准确率
test_accuracies.append(test_accuracy)
print(f"Epoch: {epoch + 1}/{epochs},Test Accuracy: {test_accuracy:.2%}")
# 如果测试集准确率提高，保存当前模型的权重
if test_accuracy > best_accuracy:
best_accuracy = test_accuracy
torch.save(model.state_dict(), best_model_path)
print(f"Best model saved with accuracy: {best_accuracy:.2%}")
print(f"Best Accuracy on test set: {best_accuracy:.2%}")
# 绘制并保存损失和准确率曲线
plt.figure(figsize=(12, 5))
# 绘制损失曲线
plt.subplot(1, 2, 1) # 选择第一个子图
plt.plot(train_losses, label='Training Loss') # 传入数据、设置标签为Training Loss
plt.xlabel('Epoch') # x轴数据
plt.ylabel('Loss')
plt.title('Training Loss over Epochs') # 设置标题
plt.legend() # 添加图例
plt.grid(True) # 添加网格
# 绘制训练集和测试集准确率曲线
plt.subplot(1, 2, 2)
plt.plot(train_accuracies, label='Train Accuracy')
plt.plot(test_accuracies, label='Test Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Train and Test Accuracy over Epochs')
plt.legend()
plt.grid(True)
# 保存图像
plt.tight_layout()
plt.savefig('loss_and_accuracy.png')
plt.show()
3 predict.py 创建
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from .model import MyNet
# 设置随机种子
torch.manual_seed(21)
# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device:{device}")
# 定义数据预处理
transform = transforms.Compose([
transforms.ToTensor(), # 将图像转为张量
transforms.Normalize((0.5,), (0.5,)) # 标准化图像数据，对于灰度图像，只需要一个通道的标准化
])
# 加载MNIST数据集
test_dataset = datasets.MNIST(root='../../datasets', train=False, download=True, transform=transform)
# 创建数据加载器
batch_size = 10
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 初始化模型并将模型移到GPU上
model = MyNet().to(device)
# 加载保存的权重
best_model_path = 'best_mnist_model.pth'
model.load_state_dict(torch.load(best_model_path))
# 设置模型为评估模式
model.eval()
# 从测试集中取10张图片
examples = enumerate(test_loader) # 从 test loader 中获取一个带索引的迭代器
# #使用 next()函数从迭代器中获取下一个批次的数据。
# next()返回一个元组，第一个元素是批次的索引(batch idx)，
# #第二个元素是一个包含输入数据和目标标签的元组(example_data，example_targets)
batch_idx, (example_data, example_targets) = next(examples)
# 将数据移动到指定的设备上(例如，CPU或GPU)，
# 以便利用设备的计算能力进行进一步的操作(如模型推理或训练)
example_data, example_targets = example_data.to(device), example_targets.to(device)
# 进行推理
with torch.no_grad(): # 使用 torch.no_grad()上下文管理器，表示在此上下文中不计算梯度。
outputs = model(example_data) # 使用模型对 example_data 进行前向传播，获取输出结果。
# 使用 torch.max()函数，获取 outputs 张量中每行的最大值及其索引。
# 1 表示在第一个维度上(通常是类别维度)寻找最大值。
_, predicted = torch.max(outputs, 1)
# 可视化并显示预测结果
fig, axes = plt.subplots(1, 10, figsize=(15, 2)) # 创建一个包含1行10列的子图(subplots)的图形，并指定每个子图的大小
for i in range(10):
axes[i].imshow(example_data[i].cpu().squeeze(), cmap='gray') # 在指定的子图(axes[i])中显示图像
axes[i].set_title(f"Pred:{predicted[i].item()}") # 在每个子图(axes[i])上设置标题，标题内容显示模型对该样本的预测结果
axes[i].axis('off') # 关闭第 i个子图(axes[i])的坐标轴显示
plt.show()
# 输出批次中部分样本的真实标签和模型的预测标签
print(f"True labels: {example_targets.cpu().numpy()}")
print(f"Predicted labels: {predicted.cpu().numpy()}")
