@浙大疏锦行
import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter import matplotlib.pyplot as plt import numpy as np import os import torchvision # 记得导入 torchvision，之前代码里用到了其功能但没导入 # 设置中文字体支持 plt.rcParams["font.family"] = ["SimHei"] plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题 # 检查GPU是否可用 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"使用设备: {device}") # 1. 数据预处理 train_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1), transforms.RandomRotation(15), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) test_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) # 2. 加载CIFAR-10数据集 train_dataset = datasets.CIFAR10( root='./data', train=True, download=True, transform=train_transform ) test_dataset = datasets.CIFAR10( root='./data', train=False, transform=test_transform ) # 3. 创建数据加载器 batch_size = 64 train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) # 4. 定义CNN模型的定义（替代原MLP） class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() # 继承父类初始化 # ---------------------- 第一个卷积块 ---------------------- # 卷积层1：输入3通道（RGB），输出32个特征图，卷积核3x3，边缘填充1像素 self.conv1 = nn.Conv2d( in_channels=3, # 输入通道数（图像的RGB通道） out_channels=32, # 输出通道数（生成32个新特征图） kernel_size=3, # 卷积核尺寸（3x3像素） padding=1 # 边缘填充1像素，保持输出尺寸与输入相同 ) # 批量归一化层：对32个输出通道进行归一化，加速训练 self.bn1 = nn.BatchNorm2d(num_features=32) # ReLU激活函数：引入非线性，公式：max(0, x) self.relu1 = nn.ReLU() # 最大池化层：窗口2x2，步长2，特征图尺寸减半（32x32→16x16） self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) # stride默认等于kernel_size # ---------------------- 第二个卷积块 ---------------------- # 卷积层2：输入32通道（来自conv1的输出），输出64通道 self.conv2 = nn.Conv2d( in_channels=32, # 输入通道数（前一层的输出通道数） out_channels=64, # 输出通道数（特征图数量翻倍） kernel_size=3, # 卷积核尺寸不变 padding=1 # 保持尺寸：16x16→16x16（卷积后）→8x8（池化后） ) self.bn2 = nn.BatchNorm2d(num_features=64) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool2d(kernel_size=2) # 尺寸减半：16x16→8x8 # ---------------------- 第三个卷积块 ---------------------- # 卷积层3：输入64通道，输出128通道 self.conv3 = nn.Conv2d( in_channels=64, # 输入通道数（前一层的输出通道数） out_channels=128, # 输出通道数（特征图数量再次翻倍） kernel_size=3, padding=1 # 保持尺寸：8x8→8x8（卷积后）→4x4（池化后） ) self.bn3 = nn.BatchNorm2d(num_features=128) self.relu3 = nn.ReLU() # 复用激活函数对象（节省内存） self.pool3 = nn.MaxPool2d(kernel_size=2) # 尺寸减半：8x8→4x4 # ---------------------- 全连接层（分类器） ---------------------- # 计算展平后的特征维度：128通道 × 4x4尺寸 = 128×16=2048维 self.fc1 = nn.Linear( in_features=128 * 4 * 4, # 输入维度（卷积层输出的特征数） out_features=512 # 输出维度（隐藏层神经元数） ) # Dropout层：训练时随机丢弃50%神经元，防止过拟合 self.dropout = nn.Dropout(p=0.5) # 输出层：将512维特征映射到10个类别（CIFAR-10的类别数） self.fc2 = nn.Linear(in_features=512, out_features=10) def forward(self, x): # 输入尺寸：[batch_size, 3, 32, 32]（batch_size=批量大小，3=通道数，32x32=图像尺寸） # ---------- 卷积块1处理 ---------- x = self.conv1(x) # 卷积后尺寸：[batch_size, 32, 32, 32]（padding=1保持尺寸） x = self.bn1(x) # 批量归一化，不改变尺寸 x = self.relu1(x) # 激活函数，不改变尺寸 x = self.pool1(x) # 池化后尺寸：[batch_size, 32, 16, 16]（32→16是因为池化窗口2x2） # ---------- 卷积块2处理 ---------- x = self.conv2(x) # 卷积后尺寸：[batch_size, 64, 16, 16]（padding=1保持尺寸） x = self.bn2(x) x = self.relu2(x) x = self.pool2(x) # 池化后尺寸：[batch_size, 64, 8, 8] # ---------- 卷积块3处理 ---------- x = self.conv3(x) # 卷积后尺寸：[batch_size, 128, 8, 8]（padding=1保持尺寸） x = self.bn3(x) x = self.relu3(x) x = self.pool3(x) # 池化后尺寸：[batch_size, 128, 4, 4] # ---------- 展平与全连接层 ---------- # 将多维特征图展平为一维向量：[batch_size, 128*4*4] = [batch_size, 2048] x = x.view(-1, 128 * 4 * 4) # -1自动计算批量维度，保持批量大小不变 x = self.fc1(x) # 全连接层：2048→512，尺寸变为[batch_size, 512] x = self.relu3(x) # 激活函数（复用relu3，与卷积块3共用） x = self.dropout(x) # Dropout随机丢弃神经元，不改变尺寸 x = self.fc2(x) # 全连接层：512→10，尺寸变为[batch_size, 10]（未激活，直接输出logits） return x # 输出未经过Softmax的logits，适用于交叉熵损失函数 # 初始化模型 model = CNN() model = model.to(device) # 将模型移至GPU（如果可用） criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer, # 指定要控制的优化器（这里是Adam） mode='min', # 监测的指标是"最小化"（如损失函数） patience=3, # 如果连续3个epoch指标没有改善，才降低LR factor=0.5, # 降低LR的比例（新LR = 旧LR × 0.5） verbose=True # 打印学习率调整信息 ) # ======================== TensorBoard 核心配置 ======================== # 创建 TensorBoard 日志目录（自动避免重复） log_dir = "runs/cifar10_cnn_exp" if os.path.exists(log_dir): version = 1 while os.path.exists(f"{log_dir}_v{version}"): version += 1 log_dir = f"{log_dir}_v{version}" writer = SummaryWriter(log_dir) # 初始化 SummaryWriter # 5. 训练模型（整合 TensorBoard 记录） def train(model, train_loader, test_loader, criterion, optimizer, scheduler, device, epochs, writer): model.train() all_iter_losses = [] iter_indices = [] global_step = 0 # 全局步骤，用于 TensorBoard 标量记录 # （可选）记录模型结构：用一个真实样本走一遍前向传播，让 TensorBoard 解析计算图 dataiter = iter(train_loader) images, labels = next(dataiter) images = images.to(device) writer.add_graph(model, images) # 写入模型结构到 TensorBoard # （可选）记录原始训练图像：可视化数据增强前/后效果 img_grid = torchvision.utils.make_grid(images[:8].cpu()) # 取前8张 writer.add_image('原始训练图像（增强前）', img_grid, global_step=0) for epoch in range(epochs): 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() # 记录迭代级损失 iter_loss = loss.item() all_iter_losses.append(iter_loss) iter_indices.append(global_step + 1) # 用 global_step 对齐 # 统计准确率 running_loss += iter_loss _, predicted = output.max(1) total += target.size(0) correct += predicted.eq(target).sum().item() # ======================== TensorBoard 标量记录 ======================== # 记录每个 batch 的损失、准确率 batch_acc = 100. * correct / total writer.add_scalar('Train/Batch Loss', iter_loss, global_step) writer.add_scalar('Train/Batch Accuracy', batch_acc, global_step) # 记录学习率（可选） writer.add_scalar('Train/Learning Rate', optimizer.param_groups[0]['lr'], global_step) # 每 200 个 batch 记录一次参数直方图（可选，耗时稍高） if (batch_idx + 1) % 200 == 0: for name, param in model.named_parameters(): writer.add_histogram(f'Weights/{name}', param, global_step) if param.grad is not None: writer.add_histogram(f'Gradients/{name}', param.grad, global_step) # 每 100 个 batch 打印控制台日志（同原代码） if (batch_idx + 1) % 100 == 0: print(f'Epoch: {epoch+1}/{epochs} | Batch: {batch_idx+1}/{len(train_loader)} ' f'| 单Batch损失: {iter_loss:.4f} | 累计平均损失: {running_loss/(batch_idx+1):.4f}') global_step += 1 # 全局步骤递增 # 计算 epoch 级训练指标 epoch_train_loss = running_loss / len(train_loader) epoch_train_acc = 100. * correct / total # ======================== TensorBoard epoch 标量记录 ======================== writer.add_scalar('Train/Epoch Loss', epoch_train_loss, epoch) writer.add_scalar('Train/Epoch Accuracy', epoch_train_acc, epoch) # 测试阶段 model.eval() test_loss = 0 correct_test = 0 total_test = 0 wrong_images = [] # 存储错误预测样本（用于可视化） wrong_labels = [] wrong_preds = [] 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_test += target.size(0) correct_test += predicted.eq(target).sum().item() # 收集错误预测样本（用于可视化） wrong_mask = (predicted != target) if wrong_mask.sum() > 0: wrong_batch_images = data[wrong_mask][:8].cpu() # 最多存8张 wrong_batch_labels = target[wrong_mask][:8].cpu() wrong_batch_preds = predicted[wrong_mask][:8].cpu() wrong_images.extend(wrong_batch_images) wrong_labels.extend(wrong_batch_labels) wrong_preds.extend(wrong_batch_preds) # 计算 epoch 级测试指标 epoch_test_loss = test_loss / len(test_loader) epoch_test_acc = 100. * correct_test / total_test # ======================== TensorBoard 测试集记录 ======================== writer.add_scalar('Test/Epoch Loss', epoch_test_loss, epoch) writer.add_scalar('Test/Epoch Accuracy', epoch_test_acc, epoch) # （可选）可视化错误预测样本 if wrong_images: wrong_img_grid = torchvision.utils.make_grid(wrong_images) writer.add_image('错误预测样本', wrong_img_grid, epoch) # 写入错误标签文本（可选） wrong_text = [f"真实: {classes[wl]}, 预测: {classes[wp]}" for wl, wp in zip(wrong_labels, wrong_preds)] writer.add_text('错误预测标签', '\n'.join(wrong_text), epoch) # 更新学习率调度器 scheduler.step(epoch_test_loss) print(f'Epoch {epoch+1}/{epochs} 完成 | 训练准确率: {epoch_train_acc:.2f}% | 测试准确率: {epoch_test_acc:.2f}%') # 关闭 TensorBoard 写入器 writer.close() # 绘制迭代级损失曲线（同原代码） plot_iter_losses(all_iter_losses, iter_indices) return epoch_test_acc # 6. 绘制迭代级损失曲线（同原代码，略） def plot_iter_losses(losses, indices): plt.figure(figsize=(10, 4)) plt.plot(indices, losses, 'b-', alpha=0.7, label='Iteration Loss') plt.xlabel('Iteration（Batch序号）') plt.ylabel('损失值') plt.title('每个 Iteration 的训练损失') plt.legend() plt.grid(True) plt.tight_layout() plt.show() # （可选）CIFAR-10 类别名 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') # 7. 执行训练（传入 TensorBoard writer） epochs = 20 print("开始使用CNN训练模型...") print(f"TensorBoard 日志目录: {log_dir}") print("训练后执行: tensorboard --logdir=runs 查看可视化") final_accuracy = train(model, train_loader, test_loader, criterion, optimizer, scheduler, device, epochs, writer) print(f"训练完成！最终测试准确率: {final_accuracy:.2f}%")
