python训练营打卡第40天

训练和测试的规范写法

知识点回顾：

1. 彩色和灰度图片测试和训练的规范写法：封装在函数中
2. 展平操作：除第一个维度batchsize外全部展平
3. dropout操作：训练阶段随机丢弃神经元，测试阶段eval模式关闭dropout

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
 
# 设备配置
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
 
# 数据预处理函数 - 处理灰度图像(MNIST)
def get_mnist_loaders(batch_size=64):
    # 灰度图像归一化
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
    ])
    
    train_dataset = datasets.MNIST(
        root='./data', 
        train=True, 
        download=True, 
        transform=transform
    )
    test_dataset = datasets.MNIST(
        root='./data', 
        train=False, 
        download=True, 
        transform=transform
    )
    
    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
 
# 数据预处理函数 - 处理彩色图像(CIFAR-10)
def get_cifar10_loaders(batch_size=64):
    # 彩色图像归一化
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # CIFAR-10的均值和标准差
    ])
    
    train_dataset = datasets.CIFAR10(
        root='./data', 
        train=True, 
        download=True, 
        transform=transform
    )
    test_dataset = datasets.CIFAR10(
        root='./data', 
        train=False, 
        download=True, 
        transform=transform
    )
    
    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
 
# 定义包含展平和dropout的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self, in_channels=1, num_classes=10):
        super(SimpleCNN, self).__init__()
        # 卷积层部分
        self.conv_layers = nn.Sequential(
            nn.Conv2d(in_channels, 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Dropout(0.25),  # 第一个dropout层
            
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Dropout(0.25)   # 第二个dropout层
        )
        
        # 展平操作后接全连接层
        self.fc_layers = nn.Sequential(
            nn.Flatten(),  # 展平操作，保持batch维度不变
            nn.Linear(64 * 7 * 7, 128),  # 假设输入尺寸为32x32，经过两次池化后为8x8，64通道
            nn.ReLU(),
            nn.Dropout(0.5),  # 全连接层后的dropout
            nn.Linear(128, num_classes)
        )
    
    def forward(self, x):
        x = self.conv_layers(x)
        x = self.fc_layers(x)
        return x
 
# 训练函数 - 规范的训练流程
def train(model, train_loader, criterion, optimizer, epoch, device):
    model.train()  # 切换到训练模式，启用dropout
    running_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        
        # 梯度清零
        optimizer.zero_grad()
        
        # 前向传播
        output = model(data)
        loss = criterion(output, target)
        
        # 反向传播和优化
        loss.backward()
        optimizer.step()
        
        # 统计训练信息
        running_loss += loss.item()
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()
        
        if batch_idx % 100 == 0:
            print(f'Epoch: {epoch} | Batch: {batch_idx} | Loss: {loss.item():.4f} | Acc: {100.*correct/total:.2f}%')
    
    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc
 
# 测试函数 - 规范的测试流程
def test(model, test_loader, criterion, device):
    model.eval()  # 切换到测试模式，关闭dropout
    test_loss = 0
    correct = 0
    total = 0
    
    with torch.no_grad():  # 不计算梯度，节省内存和计算资源
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            
            # 累加损失
            test_loss += criterion(output, target).item()
            
            # 计算准确率
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()
    
    test_loss /= len(test_loader)
    test_acc = 100. * correct / total
    print(f'Test Loss: {test_loss:.4f} | Test Accuracy: {test_acc:.2f}%')
    return test_loss, test_acc
 
# 主函数 - 整合整个流程
def main(use_color=False):
    # 选择数据集
    if use_color:
        print("Using CIFAR-10 (color images) dataset...")
        train_loader, test_loader = get_cifar10_loaders(batch_size=128)
        in_channels = 3  # 彩色图像3通道
    else:
        print("Using MNIST (grayscale images) dataset...")
        train_loader, test_loader = get_mnist_loaders(batch_size=128)
        in_channels = 1  # 灰度图像1通道
    
    # 初始化模型
    model = SimpleCNN(in_channels=in_channels, num_classes=10).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    # 训练和测试循环
    num_epochs = 5
    train_losses, train_accs, test_losses, test_accs = [], [], [], []
    
    for epoch in range(1, num_epochs + 1):
        print(f"\nEpoch {epoch}/{num_epochs}")
        train_loss, train_acc = train(model, train_loader, criterion, optimizer, epoch, device)
        test_loss, test_acc = test(model, test_loader, criterion, device)
        
        train_losses.append(train_loss)
        train_accs.append(train_acc)
        test_losses.append(test_loss)
        test_accs.append(test_acc)
    

输出结果：

Using device: cpu
Using MNIST (grayscale images) dataset...

Epoch 1/5
Epoch: 1 | Batch: 0 | Loss: 2.3310 | Acc: 7.81%
Epoch: 1 | Batch: 100 | Loss: 0.3962 | Acc: 75.91%
Epoch: 1 | Batch: 200 | Loss: 0.1773 | Acc: 84.18%
Epoch: 1 | Batch: 300 | Loss: 0.1052 | Acc: 87.58%
Epoch: 1 | Batch: 400 | Loss: 0.0723 | Acc: 89.52%
Test Loss: 0.0672 | Test Accuracy: 97.87%

Epoch 2/5
Epoch: 2 | Batch: 0 | Loss: 0.1675 | Acc: 94.53%
Epoch: 2 | Batch: 100 | Loss: 0.1031 | Acc: 96.50%
Epoch: 2 | Batch: 200 | Loss: 0.0900 | Acc: 96.50%
Epoch: 2 | Batch: 300 | Loss: 0.0752 | Acc: 96.66%
Epoch: 2 | Batch: 400 | Loss: 0.0974 | Acc: 96.68%
Test Loss: 0.0390 | Test Accuracy: 98.63%

Epoch 3/5
Epoch: 3 | Batch: 0 | Loss: 0.0644 | Acc: 98.44%
Epoch: 3 | Batch: 100 | Loss: 0.0499 | Acc: 97.41%
Epoch: 3 | Batch: 200 | Loss: 0.0949 | Acc: 97.47%
Epoch: 3 | Batch: 300 | Loss: 0.1349 | Acc: 97.53%
Epoch: 3 | Batch: 400 | Loss: 0.1658 | Acc: 97.47%
Test Loss: 0.0361 | Test Accuracy: 98.71%

Epoch 4/5
Epoch: 4 | Batch: 0 | Loss: 0.1060 | Acc: 95.31%
Epoch: 4 | Batch: 100 | Loss: 0.0539 | Acc: 97.87%
Epoch: 4 | Batch: 200 | Loss: 0.0220 | Acc: 97.69%
Epoch: 4 | Batch: 300 | Loss: 0.0484 | Acc: 97.71%
Test Loss: 0.0336 | Test Accuracy: 98.80%
Epoch 5/5
Epoch: 5 | Batch: 0 | Loss: 0.0804 | Acc: 96.88%
Epoch: 5 | Batch: 100 | Loss: 0.0255 | Acc: 98.04%
Epoch: 5 | Batch: 200 | Loss: 0.0396 | Acc: 97.97%
Epoch: 5 | Batch: 300 | Loss: 0.0869 | Acc: 97.97%
Epoch: 5 | Batch: 400 | Loss: 0.0246 | Acc: 97.98%
Test Loss: 0.0275 | Test Accuracy: 99.00%

@浙大疏锦行