代码功能

网络结构：
卷积层：
两个卷积层，每个卷积层后接 ReLU 激活函数。
最大池化层用于降低维度。
全连接层：
使用一个隐藏层（128 个神经元）和一个输出层（10 类分类任务）。
数据集：
使用 PyTorch 内置的 MNIST 数据集，其中包含手写数字的灰度图像。
训练过程：
使用交叉熵损失函数 (CrossEntropyLoss)。
优化器为 Adam，学习率设为 0.001。
每轮训练输出损失。
测试与可视化：
测试模型在测试集上的准确率。
可视化 6 张测试样本的预测结果与真实标签。

代码

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# 1. 定义卷积神经网络
class CNN(nn.Module):def __init__(self):super(CNN, self).__init__()self.conv_layers = nn.Sequential(nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),  # 卷积层 1nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2),  # 最大池化层nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),  # 卷积层 2nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层)self.fc_layers = nn.Sequential(nn.Flatten(),nn.Linear(64 * 7 * 7, 128),  # 全连接层 1nn.ReLU(),nn.Linear(128, 10)  # 全连接层 2 (10 类分类))def forward(self, x):x = self.conv_layers(x)x = self.fc_layers(x)return x# 2. 加载 MNIST 数据集
def load_data(batch_size=64):transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))  # 标准化])train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)return train_loader, test_loader# 3. 训练 CNN
def train_cnn(model, train_loader, criterion, optimizer, epochs=5):model.train()for epoch in range(epochs):total_loss = 0for images, labels in train_loader:images, labels = images.cuda(), labels.cuda()  # 将数据移动到 GPU（如适用）# 前向传播outputs = model(images)loss = criterion(outputs, labels)# 反向传播与优化optimizer.zero_grad()loss.backward()optimizer.step()total_loss += loss.item()print(f"Epoch [{epoch + 1}/{epochs}], Loss: {total_loss / len(train_loader):.4f}")# 4. 测试 CNN
def test_cnn(model, test_loader):model.eval()correct = 0total = 0with torch.no_grad():for images, labels in test_loader:images, labels = images.cuda(), labels.cuda()  # 将数据移动到 GPU（如适用）outputs = model(images)_, predicted = torch.max(outputs, 1)total += labels.size(0)correct += (predicted == labels).sum().item()print(f"Test Accuracy: {100 * correct / total:.2f}%")# 5. 可视化测试结果
def visualize_predictions(model, test_loader):model.eval()images, labels = next(iter(test_loader))images, labels = images.cuda(), labels.cuda()outputs = model(images)_, predicted = torch.max(outputs, 1)# 绘制图像与预测结果images, labels, predicted = images.cpu(), labels.cpu(), predicted.cpu()plt.figure(figsize=(12, 8))for i in range(6):plt.subplot(2, 3, i + 1)plt.imshow(images[i].squeeze(), cmap='gray')plt.title(f"True: {labels[i]}, Pred: {predicted[i]}")plt.axis('off')plt.show()# 主程序
if __name__ == "__main__":# 加载数据train_loader, test_loader = load_data()# 初始化网络、损失函数和优化器model = CNN().cuda()  # 将模型移动到 GPU（如适用）criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练和测试模型train_cnn(model, train_loader, criterion, optimizer, epochs=5)test_cnn(model, test_loader)# 可视化部分测试结果visualize_predictions(model, test_loader)