AI大模型研发设计实战应用指南-优快云博客

PyTorch 实现一个 CNN 实例

以下示例展示如何用 PyTorch 构建一个简单的 CNN 模型，用于 MNIST 数据集的数字分类。

主要步骤：

数据加载与预处理：使用 torchvision 加载和预处理 MNIST 数据。
模型构建：定义卷积层、池化层和全连接层。
训练：通过损失函数和优化器进行模型训练。
评估：测试集上计算模型的准确率。
可视化：展示部分测试样本及其预测结果。

Python

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

Python

# 1. 数据加载与预处理
transform = transforms.Compose([
transforms.ToTensor(), # 转为张量
transforms.Normalize((0.5,), (0.5,)) # 归一化到 [-1, 1]
])

Python

# 加载 MNIST 数据集
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 = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

Python

# 2. 定义 CNN 模型
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
# 定义卷积层
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
# 定义全连接层
self.fc1 = nn.Linear(64 * 7 * 7, 128) # 展平后输入到全连接层
self.fc2 = nn.Linear(128, 10) # 10 个类别
def forward(self, x):
x = F.relu(self.conv1(x)) # 第一层卷积 + ReLU
x = F.max_pool2d(x, 2) # 最大池化
x = F.relu(self.conv2(x)) # 第二层卷积 + ReLU
x = F.max_pool2d(x, 2) # 最大池化
x = x.view(-1, 64 * 7 * 7) # 展平
x = F.relu(self.fc1(x)) # 全连接层 + ReLU
x = self.fc2(x) # 最后一层输出
return x
# 创建模型实例
model = SimpleCNN()

Python

# 3. 定义损失函数与优化器
criterion = nn.CrossEntropyLoss() # 多分类交叉熵损失
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

Python

# 4. 模型训练
num_epochs = 5
model.train() # 设置模型为训练模式
for epoch in range(num_epochs):
total_loss = 0
for images, labels in train_loader:
outputs = model(images) # 前向传播
loss = criterion(outputs, labels) # 计算损失
optimizer.zero_grad() # 清空梯度
loss.backward() # 反向传播
optimizer.step() # 更新参数
total_loss += loss.item()
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss / len(train_loader):.4f}")

Python

# 5. 模型测试
model.eval() # 设置模型为评估模式
correct = 0
total = 0
with torch.no_grad(): # 关闭梯度计算
for images, labels in test_loader:
outputs = model(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy:.2f}%")