深度学习入门：CNN 图像分类实战

# 安装 PyTorch（根据你的 CUDA 版本选择） pip install torch torchvision torchaudio # 如果使用 CPU 版本 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu # 验证安装 python -c "import torch; print(torch.__version__)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())"

import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms import matplotlib.pyplot as plt import numpy as np

class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() # 卷积层 1 # 输入：[batch_size, 1, 28, 28] # 输出：[batch_size, 32, 28, 28] self.conv1 = nn.Sequential( nn.Conv2d( in_channels=1, # 输入通道数（灰度图） out_channels=32, # 输出通道数（滤波器数量） kernel_size=3, # 卷积核大小 stride=1, # 步长 padding=1 # 填充，保持尺寸不变 ), nn.ReLU(), # 激活函数 nn.MaxPool2d(2) # 池化层，尺寸减半 → [batch_size, 32, 14, 14] ) # 卷积层 2 # 输入：[batch_size, 32, 14, 14] # 输出：[batch_size, 64, 14, 14] → 池化后 [batch_size, 64, 7, 7] self.conv2 = nn.Sequential( nn.Conv2d(32, 64, 3, 1, 1), nn.ReLU(), nn.MaxPool2d(2) ) # 全连接层 # 输入：[batch_size, 64*7*7] # 输出：[batch_size, 10] self.fc = nn.Sequential( nn.Linear(64 * 7 * 7, 128), # 隐藏层 nn.ReLU(), nn.Dropout(0.5), # Dropout 防止过拟合 nn.Linear(128, 10) # 输出层（10 个数字类别） ) def forward(self, x): # 卷积和池化 x = self.conv1(x) x = self.conv2(x) # 展平，准备进入全连接层 x = x.view(x.size(0), -1) # [batch_size, 64*7*7] # 全连接层 output = self.fc(x) return output

# 定义数据转换 transform = transforms.Compose([ transforms.ToTensor(), # 转换为 Tensor transforms.Normalize((0.1307,), (0.3081,)) # 标准化 ]) # 下载并加载训练集 train_dataset = datasets.MNIST( root='./data', train=True, download=True, transform=transform ) train_loader = DataLoader( dataset=train_dataset, batch_size=64, shuffle=True # 打乱数据 ) # 加载测试集 test_dataset = datasets.MNIST( root='./data', train=False, download=True, transform=transform ) test_loader = DataLoader( dataset=test_dataset, batch_size=1000, shuffle=False )

# 初始化模型、损失函数和优化器 model = CNN() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # 训练配置 EPOCHS = 5 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model.to(DEVICE) def train_epoch(model, loader, optimizer, criterion, device): model.train() total_loss = 0 correct = 0 for batch_idx, (data, target) in enumerate(loader): data, target = data.to(device), target.to(device) # 前向传播 output = model(data) loss = criterion(output, target) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() avg_loss = total_loss / len(loader) accuracy = 100. * correct / len(loader.dataset) return avg_loss, accuracy # 开始训练 for epoch in range(1, EPOCHS + 1): loss, acc = train_epoch(model, train_loader, optimizer, criterion, DEVICE) print(f'Epoch {epoch}: Loss={loss:.4f}, Accuracy={acc:.2f}%')

def evaluate(model, loader, criterion, device): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += criterion(output, target).item() pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() avg_loss = test_loss / len(loader) accuracy = 100. * correct / len(loader.dataset) return avg_loss, accuracy # 在测试集上评估 test_loss, test_acc = evaluate(model, test_loader, criterion, DEVICE) print(f'\nTest Results: Loss={test_loss:.4f}, Accuracy={test_acc:.2f}%')

def visualize_predictions(model, dataset, num_samples=10): model.eval() fig, axes = plt.subplots(2, 5, figsize=(12, 5)) axes = axes.flatten() for i in range(num_samples): img, label = dataset[i] # 添加批次维度 img_tensor = img.unsqueeze(0).to(DEVICE) # 预测 with torch.no_grad(): output = model(img_tensor) prediction = output.argmax(dim=1).item() # 显示图像 axes[i].imshow(img.squeeze(), cmap='gray') axes[i].set_title(f'Pred: {prediction}\nTrue: {label}') axes[i].axis('off') plt.tight_layout() plt.savefig('predictions.png', dpi=300) plt.show() # 可视化预测结果 visualize_predictions(model, test_dataset)

# 使用学习率调度器 scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=10, # 每 10 个 epoch 调整一次 gamma=0.1 # 学习率衰减为原来的 0.1 倍 ) # 在每个 epoch 结束后调用 scheduler.step()

# 增加数据变换，提升泛化能力 transform_augmented = transforms.Compose([ transforms.RandomRotation(10), # 随机旋转 ±10 度 transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ])

# 在卷积层后添加 BN 层 self.conv1 = nn.Sequential( nn.Conv2d(1, 32, 3, 1, 1), nn.BatchNorm2d(32), # 加速收敛，提高稳定性 nn.ReLU(), nn.MaxPool2d(2) )