DAY45 Tensorboard使用介绍

@浙大疏锦行

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
import numpy as np
import os
from datetime import datetime

# 设置随机种子保证可重复性
torch.manual_seed(42)
np.random.seed(42)

# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 数据预处理和增强
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

# 加载CIFAR-10数据集
train_dataset = datasets.CIFAR10(
    root='./data', 
    train=True, 
    download=True, 
    transform=transform_train
)
test_dataset = datasets.CIFAR10(
    root='./data', 
    train=False, 
    download=True, 
    transform=transform_test
)

train_loader = DataLoader(
    train_dataset, 
    batch_size=128, 
    shuffle=True, 
    num_workers=2
)
test_loader = DataLoader(
    test_dataset, 
    batch_size=100, 
    shuffle=False, 
    num_workers=2
)

# 类别名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

# 创建TensorBoard的SummaryWriter
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
log_dir = f'runs/resnet18_cifar10_{timestamp}'
writer = SummaryWriter(log_dir=log_dir)
print(f"TensorBoard日志目录: {log_dir}")

# 创建ResNet18模型（使用预训练权重）
def create_resnet18_finetune(num_classes=10, freeze_backbone=True):
    """
    创建用于微调的ResNet18模型
    
    Args:
        num_classes: 输出类别数
        freeze_backbone: 是否冻结主干网络（前几个epoch可以冻结，之后解冻）
    """
    # 加载预训练的ResNet18
    model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
    
    # 冻结所有卷积层（在微调开始时）
    if freeze_backbone:
        for param in model.parameters():
            param.requires_grad = False
    
    # 修改最后的全连接层以适应CIFAR-10
    num_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Dropout(0.5),
        nn.Linear(num_features, 256),
        nn.ReLU(),
        nn.Dropout(0.3),
        nn.Linear(256, num_classes)
    )
    
    return model

# 训练函数
def train_epoch(model, train_loader, criterion, optimizer, epoch, device, writer=None):
    """训练一个epoch"""
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)
        
        # 前向传播
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
        # 反向传播
        loss.backward()
        optimizer.step()
        
        # 统计
        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        
        # 每100个batch打印一次进度
        if (batch_idx + 1) % 100 == 0:
            batch_loss = loss.item()
            batch_acc = 100. * predicted.eq(targets).sum().item() / targets.size(0)
            print(f'Epoch: {epoch+1}, Batch: {batch_idx+1}/{len(train_loader)}, '
                  f'Loss: {batch_loss:.4f}, Acc: {batch_acc:.2f}%')
            
            # 记录每个batch的学习率
            if writer:
                writer.add_scalar('Learning Rate', optimizer.param_groups[0]['lr'], 
                                epoch * len(train_loader) + batch_idx)
    
    # 计算整个epoch的指标
    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100. * correct / total
    
    # 记录到TensorBoard
    if writer:
        writer.add_scalar('Loss/Train', epoch_loss, epoch)
        writer.add_scalar('Accuracy/Train', epoch_acc, epoch)
    
    return epoch_loss, epoch_acc

# 测试函数
def test_epoch(model, test_loader, criterion, epoch, device, writer=None):
    """测试模型"""
    model.eval()
    test_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, targets in test_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    
    # 计算整体指标
    test_loss = test_loss / len(test_loader)
    test_acc = 100. * correct / total
    
    # 记录到TensorBoard
    if writer:
        writer.add_scalar('Loss/Test', test_loss, epoch)
        writer.add_scalar('Accuracy/Test', test_acc, epoch)
    
    return test_loss, test_acc

# 主训练循环
def train_resnet18_finetune(epochs=30, lr=0.001, unfreeze_epoch=5):
    """
    微调ResNet18
    
    Args:
        epochs: 总训练轮数
        lr: 学习率
        unfreeze_epoch: 解冻主干网络的epoch
    """
    print("=" * 50)
    print("开始微调ResNet18")
    print("=" * 50)
    
    # 创建模型
    model = create_resnet18_finetune(num_classes=10, freeze_backbone=True)
    model = model.to(device)
    
    # 打印模型参数
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"总参数: {total_params:,}")
    print(f"可训练参数: {trainable_params:,}")
    
    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    
    # 第一阶段：只训练最后的全连接层
    optimizer = optim.AdamW(model.fc.parameters(), lr=lr, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    
    best_acc = 0.0
    
    for epoch in range(epochs):
        print(f"\nEpoch {epoch+1}/{epochs}")
        print("-" * 30)
        
        # 在指定epoch解冻主干网络
        if epoch == unfreeze_epoch:
            print("解冻主干网络，开始微调所有层...")
            for param in model.parameters():
                param.requires_grad = True
            
            # 重新定义优化器，包含所有参数
            optimizer = optim.AdamW(
                model.parameters(), 
                lr=lr/10,  # 解冻后使用更小的学习率
                weight_decay=1e-4
            )
            scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs-unfreeze_epoch)
            
            # 更新可训练参数计数
            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
            print(f"解冻后可训练参数: {trainable_params:,}")
        
        # 训练
        train_loss, train_acc = train_epoch(
            model, train_loader, criterion, optimizer, 
            epoch, device, writer
        )
        
        # 测试
        test_loss, test_acc = test_epoch(
            model, test_loader, criterion, epoch, 
            device, writer
        )
        
        # 学习率调度
        scheduler.step()
        
        # 打印结果
        print(f"训练结果: Loss: {train_loss:.4f}, Acc: {train_acc:.2f}%")
        print(f"测试结果: Loss: {test_loss:.4f}, Acc: {test_acc:.2f}%")
        print(f"当前学习率: {optimizer.param_groups[0]['lr']:.6f}")
        
        # 保存最佳模型
        if test_acc > best_acc:
            best_acc = test_acc
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'test_acc': test_acc,
                'train_acc': train_acc,
            }, f'best_resnet18_cifar10.pth')
            print(f"保存最佳模型，准确率: {test_acc:.2f}%")
    
    print("\n" + "=" * 50)
    print(f"训练完成！最佳测试准确率: {best_acc:.2f}%")
    print("=" * 50)
    
    return model, best_acc

# 可视化函数
def visualize_results(model, test_loader, writer, num_images=8):
    """在TensorBoard中可视化一些测试图片和预测结果"""
    model.eval()
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    
    # 只取前num_images张图片
    images = images[:num_images]
    labels = labels[:num_images]
    images = images.to(device)
    
    with torch.no_grad():
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
    
    # 将图片添加到TensorBoard
    # 反归一化用于显示
    mean = torch.tensor([0.4914, 0.4822, 0.4465]).view(3, 1, 1)
    std = torch.tensor([0.2023, 0.1994, 0.2010]).view(3, 1, 1)
    images_denorm = images.cpu() * std + mean
    images_denorm = torch.clamp(images_denorm, 0, 1)
    
    # 创建图片网格
    img_grid = make_grid(images_denorm, nrow=4, normalize=False)
    
    # 添加到TensorBoard
    writer.add_image('Test Images/Predictions', img_grid, 0)
    
    # 添加预测文本
    print("\n预测示例:")
    for i in range(min(num_images, len(images))):
        pred_label = predicted[i].item()
        true_label = labels[i].item()
        is_correct = pred_label == true_label
        status = "✓" if is_correct else "✗"
        print(f"图片{i+1}: 真实={classes[true_label]:8s}, 预测={classes[pred_label]:8s} {status}")
    
    return images, labels, predicted

# 运行训练
if __name__ == "__main__":
    # 创建日志目录
    os.makedirs('models', exist_ok=True)
    
    # 开始训练
    trained_model, best_acc = train_resnet18_finetune(
        epochs=30,  # 可以调整epoch数量
        lr=0.001,   # 初始学习率
        unfreeze_epoch=5  # 在第5个epoch解冻主干网络
    )
    
    # 可视化一些结果
    print("\n可视化预测结果...")
    images, labels, predictions = visualize_results(trained_model, test_loader, writer, num_images=8)
    
    # 记录模型图结构
    print("记录模型图结构到TensorBoard...")
    dummy_input = torch.randn(1, 3, 32, 32).to(device)
    writer.add_graph(trained_model, dummy_input)
    
    # 记录超参数
    print("记录超参数到TensorBoard...")
    writer.add_hparams(
        {
            'lr': 0.001,
            'batch_size': 128,
            'epochs': 30,
            'unfreeze_epoch': 5,
            'weight_decay': 1e-4
        },
        {
            'best_accuracy': best_acc,
            'final_train_acc': 0,  # 这里可以记录实际的训练准确率
            'final_test_acc': 0    # 这里可以记录实际的测试准确率
        }
    )
    
    # 关闭TensorBoard写入器
    writer.close()
    
    print(f"\nTensorBoard日志保存在: {log_dir}")
    print("在终端运行以下命令查看结果:")
    print(f"tensorboard --logdir={log_dir}")
    print("\n或者指定端口:")
    print(f"tensorboard --logdir={log_dir} --port=6006")
    
    # 加载最佳模型并测试
    print("\n加载最佳模型进行最终测试...")
    checkpoint = torch.load('best_resnet18_cifar10.pth')
    trained_model.load_state_dict(checkpoint['model_state_dict'])
    
    # 最终测试
    test_loss, test_acc = test_epoch(
        trained_model, test_loader, 
        nn.CrossEntropyLoss(), 0, device
    )
    print(f"最佳模型测试准确率: {test_acc:.2f}%")