01.手写数字识别

"""
CIFAR-10图像分类训练脚本
优化内容：代码结构分层、注释增强、模型结构验证、功能模块化
"""

import os
import time
from typing import Tuple

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from torchvision import transforms


# --------------------------
# 设备配置
# --------------------------
def configure_device() -> torch.device:
    """自动选择可用设备并打印信息"""
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"使用CUDA加速: {torch.cuda.get_device_name(0)}")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print("使用Apple Metal加速 (MPS)")
    else:
        device = torch.device("cpu")
        print("使用CPU进行训练 (性能可能受限)")
    return device


DEVICE = configure_device()
os.makedirs("checkpoints", exist_ok=True)  # 创建模型保存目录


# --------------------------
# 1.data 数据准备
# --------------------------
def prepare_dataloaders(
        batch_size: int = 128,
        num_workers: int = 2
) -> Tuple[DataLoader, DataLoader]:
    """
    准备CIFAR-10数据集加载器
    参数：
        batch_size: 训练批次大小
        num_workers: 数据加载并行进程数
    返回：
        (训练集加载器, 测试集加载器)
    """
    # 数据增强配置 (仅训练集)
    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),  # 随机裁剪增加位置鲁棒性
        transforms.RandomHorizontalFlip(),  # 水平翻转增加多样性
        transforms.ToTensor(),
        transforms.Normalize(
            mean=(0.4914, 0.4822, 0.4465),  # CIFAR-10通道均值
            std=(0.2023, 0.1994, 0.2010)  # CIFAR-10通道标准差
        )
    ])

    # 测试集转换 (无数据增强)
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(
            mean=(0.4914, 0.4822, 0.4465),
            std=(0.2023, 0.1994, 0.2010)
        )
    ])

    # 加载数据集
    train_set = torchvision.datasets.CIFAR10(
        root="./data",
        train=True,
        download=True,
        transform=train_transform
    )
    test_set = torchvision.datasets.CIFAR10(
        root="./data",
        train=False,
        download=True,
        transform=test_transform
    )

    # 数据集 (DataLoader作用是多轮循环从数据集中加载数据)
    train_loader = DataLoader(
        train_set,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True  # 加速数据传到GPU
    )
    # 测试集
    test_loader = DataLoader(
        test_set,
        batch_size=100,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )

    return train_loader, test_loader


# --------------------------
# 2.net 模型架构
# --------------------------
class CIFAR10Model(nn.Module):
    """
    CIFAR-10分类模型架构
    结构说明：
    - 3个卷积块 (Conv2d + ReLU + MaxPool)
    - 全连接层 + Dropout正则化
    - 输入尺寸: 3x32x32 (RGB图像)
    - 输出尺寸: 10 (CIFAR-10类别数)
    """

    def __init__(self):
        super().__init__()
        self.feature_extractor = nn.Sequential(
            # 卷积块1: 输入3通道，输出32通道
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),  # 输出尺寸: 32x16x16

            # 卷积块2: 输入32通道，输出64通道
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),  # 输出尺寸: 64x8x8

            # 卷积块3: 输入64通道，输出128通道
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)  # 输出尺寸: 128x4x4
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256),  # 全连接层
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),  # 防止过拟合
            nn.Linear(256, 10)  # 输出层
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        features = self.feature_extractor(x)
        return self.classifier(features)


# --------------------------
# 训练组件
# --------------------------
def create_model(device: torch.device) -> nn.Module:
    """初始化模型并转移到指定设备"""
    model = CIFAR10Model().to(device)
    print("模型架构：")
    print(model)  # 打印模型结构
    return model


def evaluate_model(
        model: nn.Module,
        data_loader: DataLoader,
        device: torch.device
) -> float:
    """
    模型评估函数
    返回：
        准确率百分比
    """
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in data_loader:
            images = images.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)

            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)

            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    return 100.0 * correct / total


# --------------------------
# 主训练流程
# --------------------------
def main():
    # 初始化组件
    train_loader, test_loader = prepare_dataloaders()
    model = create_model(DEVICE)

    # 3.loss损失函数
    criterion = nn.CrossEntropyLoss()
    # 4.optimizer
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = StepLR(optimizer, step_size=10, gamma=0.1)  # 每10个epoch学习率下降10倍
    best_acc = 0.0  # 记录最佳准确率

    # 训练循环
    print("\n开始训练...")
    for epoch in range(1, 21):  # 训练20个epoch
        epoch_start = time.time()

        # 5.训练阶段
        model.train()
        total_loss = 0.0

        for batch_idx, (images, labels) in enumerate(train_loader):
            images = images.to(DEVICE, non_blocking=True)
            labels = labels.to(DEVICE, non_blocking=True)

            # 前向传播
            outputs = model(images)
            loss = criterion(outputs, labels)

            # 反向传播与优化
            optimizer.zero_grad(set_to_none=True)  # 更高效的梯度清零
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

            # 每100批次打印进度
            if (batch_idx + 1) % 100 == 0:
                print(f"Epoch [{epoch}/20] | Batch [{batch_idx + 1}/{len(train_loader)}] "
                      f"| Loss: {loss.item():.4f}")

        # 更新学习率
        scheduler.step()

        # 6.验证阶段
        avg_loss = total_loss / len(train_loader)
        val_acc = evaluate_model(model, test_loader, DEVICE)

        # 保存最佳模型
        if val_acc > best_acc:
            print(f"准确率提升 {best_acc:.2f}% → {val_acc:.2f}%，保存模型中...")
            torch.save(model.state_dict(), "./checkpoints/best_model.pth")
            best_acc = val_acc

        # 打印epoch摘要
        epoch_time = time.time() - epoch_start
        print(f"Epoch [{epoch}/20] 耗时: {epoch_time:.1f}s | "
              f"平均损失: {avg_loss:.4f} | 验证准确率: {val_acc:.2f}%")
        print("-" * 60)

    # 最终评估
    print("\n训练完成，加载最佳模型进行测试...")
    model.load_state_dict(torch.load("./checkpoints/best_model.pth"))
    final_acc = evaluate_model(model, test_loader, DEVICE)
    print(f"最终测试准确率: {final_acc:.2f}%")


if __name__ == "__main__":
    main()