import numpy as np
import torch
import torchvision
from torchvision.datasets import mnist
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim
from torch import nn
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
==================== 定义超参数 ====================
train_batch_size = 64 # 训练时每个批次的样本数量
test_batch_size = 128 # 测试时每个批次的样本数量
learning_rate = 0.01 # 初始学习率
num_epoches = 20 # 训练的总轮数
momentum = 0.9 # SGD优化器的动量参数
==================== 数据准备和预处理 ====================
定义预处理函数
transform = transforms.Compose([
transforms.ToTensor(), # 将PIL图像转换为Tensor，并自动归一化到[0,1]范围
transforms.Normalize([0.5], [0.5]) # 标准化到[-1,1]范围，公式：(x-0.5)/0.5
])
下载数据，并对数据进行预处理
train_dataset = mnist.MNIST('../data/', train=True, transform=transform, download=True)
训练数据集：存储在../data/目录，训练模式，应用预处理变换，如果不存在则下载
test_dataset = mnist.MNIST('../data/', train=False, transform=transform)
测试数据集：存储在../data/目录，测试模式，应用相同的预处理变换
得到数据加载器
train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
训练数据加载器：批次大小64，打乱数据顺序
test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False)
测试数据加载器：批次大小128，不打乱数据顺序
==================== 可视化源数据 ====================
获取一批测试数据用于可视化
examples = enumerate(test_loader) # 枚举测试数据加载器，返回索引和批次
batch_idx, (example_data, example_targets) = next(examples) # 获取第一个批次
print(f"示例数据形状: {example_data.shape}") # 打印数据形状，应该是 [128, 1, 28, 28]
显示前6个图像
fig = plt.figure(figsize=(10, 6)) # 创建图形，尺寸10x6英寸
for i in range(6):
plt.subplot(2, 3, i+1) # 创建2行3列的子图，当前是第i+1个
plt.tight_layout() # 自动调整子图参数，使之填充整个图像区域
plt.imshow(example_data[i][0], cmap='gray', interpolation='none') # 显示灰度图像
plt.title('真实标签: {}'.format(example_targets[i])) # 设置标题为真实标签
plt.xticks([]) # 移除x轴刻度
plt.yticks([]) # 移除y轴刻度
plt.savefig('mnist_samples.png') # 保存图像到文件
print("MNIST样本图像已保存到: mnist_samples.png")
==================== 构建模型 ====================
class Net(nn.Module): # 定义神经网络类，继承自nn.Module
def init(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
super(Net, self).init() # 调用父类构造函数
self.flatten = nn.Flatten() # 展平层，将多维输入一维化
使用Sequential构建网络层
self.layer1 = nn.Sequential( # 第一个隐藏层序列
nn.Linear(in_dim, n_hidden_1), # 全连接层，输入维度in_dim，输出维度n_hidden_1
nn.BatchNorm1d(n_hidden_1) # 批归一化层，加速训练并提高稳定性
)
self.layer2 = nn.Sequential( # 第二个隐藏层序列
nn.Linear(n_hidden_1, n_hidden_2), # 全连接层
nn.BatchNorm1d(n_hidden_2) # 批归一化层
)
self.out = nn.Sequential( # 输出层序列
nn.Linear(n_hidden_2, out_dim) # 输出层，输出维度out_dim（10个类别）
)
def forward(self, x): # 定义前向传播过程
x = self.flatten(x) # 展平输入，从[batch,1,28,28]变为[batch,784]
x = F.relu(self.layer1(x)) # 第一层：线性变换+批归一化+ReLU激活
x = F.relu(self.layer2(x)) # 第二层：线性变换+批归一化+ReLU激活
x = F.softmax(self.out(x), dim=1) # 输出层：线性变换+softmax，按行计算概率分布
return x
==================== 模型实例化和配置 ====================
设置设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
如果有GPU则使用GPU，否则使用CPU
print(f"使用设备: {device}")
实例化模型
model = Net(28 * 28, 300, 100, 10) # 输入28x28=784，隐藏层300和100，输出10个类别
model.to(device) # 将模型移动到指定设备（GPU或CPU）
定义损失函数和优化器
criterion = nn.CrossEntropyLoss() # 交叉熵损失函数，适用于多分类问题
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)
SGD优化器：传入模型参数、学习率、动量
打印模型结构
print("模型结构:")
print(model)
==================== 训练模型 ====================
def main(): # 定义主函数，封装训练流程
初始化记录列表
losses = [] # 记录每个epoch的训练损失
acces = [] # 记录每个epoch的训练准确率
eval_losses = [] # 记录每个epoch的测试损失
eval_acces = [] # 记录每个epoch的测试准确率
创建TensorBoard写入器
writer = SummaryWriter(log_dir='logs', comment='train-loss')
创建SummaryWriter对象，日志目录为'logs'，注释为'train-loss'
print("开始训练...")
for epoch in range(num_epoches): # 遍历每个训练轮次
train_loss = 0 # 初始化当前epoch的训练损失
train_acc = 0 # 初始化当前epoch的训练准确率
设置为训练模式
model.train() # 启用dropout和batch normalization的训练模式
动态修改参数学习率
if epoch % 5 == 0: # 每5个epoch调整一次学习率
optimizer.param_groups[0]['lr'] *= 0.9 # 学习率乘以0.9（衰减）
print("学习率: {:.6f}".format(optimizer.param_groups[0]['lr']))
训练阶段
for img, label in train_loader: # 遍历训练数据加载器中的每个批次
img = img.to(device) # 将图像数据移动到指定设备
label = label.to(device) # 将标签数据移动到指定设备
前向传播
out = model(img) # 模型预测输出
loss = criterion(out, label) # 计算损失
反向传播
optimizer.zero_grad() # 清空之前的梯度
loss.backward() # 反向传播计算梯度
optimizer.step() # 更新模型参数
记录误差
train_loss += loss.item() # 累加批次损失
计算分类的准确率
_, pred = out.max(1) # 获取预测类别（最大概率的索引）
num_correct = (pred == label).sum().item() # 计算正确预测的数量
acc = num_correct / img.shape[0] # 计算准确率
train_acc += acc # 累加批次准确率
记录训练损失和准确率
avg_train_loss = train_loss / len(train_loader) # 计算平均训练损失
avg_train_acc = train_acc / len(train_loader) # 计算平均训练准确率
losses.append(avg_train_loss) # 添加到训练损失列表
acces.append(avg_train_acc) # 添加到训练准确率列表
保存到TensorBoard
writer.add_scalar('Train/Loss', avg_train_loss, epoch) # 记录训练损失
writer.add_scalar('Train/Accuracy', avg_train_acc, epoch) # 记录训练准确率
在测试集上检验结果
eval_loss = 0 # 初始化当前epoch的测试损失
eval_acc = 0 # 初始化当前epoch的测试准确率
将模型改为预测模式
model.eval() # 禁用dropout和使用训练阶段的batch normalization统计量
with torch.no_grad(): # 测试阶段不需要计算梯度，节省内存和计算资源
for img, label in test_loader: # 遍历测试数据加载器中的每个批次
img = img.to(device) # 将图像数据移动到指定设备
label = label.to(device) # 将标签数据移动到指定设备
前向传播
out = model(img) # 模型预测输出
loss = criterion(out, label) # 计算损失
记录误差
eval_loss += loss.item() # 累加批次损失
记录准确率
_, pred = out.max(1) # 获取预测类别
num_correct = (pred == label).sum().item() # 计算正确预测的数量
acc = num_correct / img.shape[0] # 计算准确率
eval_acc += acc # 累加批次准确率
记录测试损失和准确率
avg_eval_loss = eval_loss / len(test_loader) # 计算平均测试损失
avg_eval_acc = eval_acc / len(test_loader) # 计算平均测试准确率
eval_losses.append(avg_eval_loss) # 添加到测试损失列表
eval_acces.append(avg_eval_acc) # 添加到测试准确率列表
保存到TensorBoard
writer.add_scalar('Test/Loss', avg_eval_loss, epoch) # 记录测试损失
writer.add_scalar('Test/Accuracy', avg_eval_acc, epoch) # 记录测试准确率
打印每个epoch的结果
print('epoch: {}, Train Loss: {:.4f}, Train Acc: {:.4f}, Test Loss: {:.4f}, Test Acc: {:.4f}'
.format(epoch, avg_train_loss, avg_train_acc, avg_eval_loss, avg_eval_acc))
关闭TensorBoard写入器
writer.close()
print('训练完成!')
==================== 可视化训练结果 ====================
绘制训练损失
plt.figure(figsize=(12, 4)) # 创建图形，尺寸12x4英寸
plt.subplot(1, 2, 1) # 创建1行2列的子图，当前是第1个
plt.title('训练损失') # 设置子图标题
plt.plot(np.arange(len(losses)), losses, label='Train Loss') # 绘制训练损失曲线
plt.plot(np.arange(len(eval_losses)), eval_losses, label='Test Loss') # 绘制测试损失曲线
plt.xlabel('Epoch') # 设置x轴标签
plt.ylabel('Loss') # 设置y轴标签
plt.legend() # 显示图例
plt.subplot(1, 2, 2) # 创建1行2列的子图，当前是第2个
plt.title('准确率') # 设置子图标题
plt.plot(np.arange(len(acces)), acces, label='Train Accuracy') # 绘制训练准确率曲线
plt.plot(np.arange(len(eval_acces)), eval_acces, label='Test Accuracy') # 绘制测试准确率曲线
plt.xlabel('Epoch') # 设置x轴标签
plt.ylabel('Accuracy') # 设置y轴标签
plt.legend() # 显示图例
plt.tight_layout() # 自动调整子图参数
plt.savefig('training_results.png') # 保存训练结果图
print("训练结果图已保存到: training_results.png")
==================== 最终测试 ====================
model.eval() # 设置为评估模式
final_test_acc = 0 # 初始化最终测试准确数
total_samples = 0 # 初始化总样本数
with torch.no_grad(): # 不计算梯度
for img, label in test_loader: # 遍历测试集
img = img.to(device) # 移动数据到设备
label = label.to(device) # 移动标签到设备
out = model(img) # 模型预测
_, pred = out.max(1) # 获取预测结果
final_test_acc += (pred == label).sum().item() # 累加正确预测数
total_samples += label.size(0) # 累加总样本数
final_accuracy = 100 * final_test_acc / total_samples # 计算最终准确率百分比
print(f'\n最终测试准确率: {final_accuracy:.2f}%') # 打印最终准确率
==================== 显示一些预测结果 ====================
model.eval() # 设置为评估模式
with torch.no_grad():
获取一批测试数据
dataiter = iter(test_loader) # 创建测试数据迭代器
images, labels = next(dataiter) # 获取下一个批次
images, labels = images.to(device), labels.to(device) # 移动数据到设备
进行预测
outputs = model(images) # 模型预测
_, predicted = torch.max(outputs, 1) # 获取预测类别
显示前12个预测结果
fig, axes = plt.subplots(3, 4, figsize=(12, 9)) # 创建3行4列的子图
axes = axes.ravel() # 将子图数组展平为一维
for i in range(12): # 遍历前12个样本
axes[i].imshow(images[i].cpu().numpy()[0], cmap='gray') # 显示图像
axes[i].set_title(f'预测: {predicted[i].item()}, 真实: {labels[i].item()}') # 设置标题
axes[i].axis('off') # 关闭坐标轴
如果预测错误，用红色标题突出显示
if predicted[i] != labels[i]:
axes[i].title.set_color('red') # 预测错误时标题变为红色
plt.tight_layout() # 自动调整布局
plt.savefig('prediction_results.png') # 保存预测结果图
print("预测结果图已保存到: prediction_results.png")
使用if name == 'main'保护主程序入口
if name == 'main':
设置matplotlib后端为Agg，避免显示问题（不显示图形窗口，只保存到文件）
import matplotlib
matplotlib.use('Agg')
main() # 调用主函数开始训练