作业：对resnet18在cifar10上采用微调策略下，用tensorboard监控训练过程。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.models import resnet18
from torch.utils.tensorboard import SummaryWriter
import time
import os

# 设置随机种子确保结果可复现
torch.manual_seed(42)

# 数据预处理
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数据集
trainset = datasets.CIFAR10(root='./data', train=True,
                             download=True, transform=transform_train)
trainloader = DataLoader(trainset, batch_size=128,
                          shuffle=True, num_workers=2)

testset = datasets.CIFAR10(root='./data', train=False,
                            download=True, transform=transform_test)
testloader = DataLoader(testset, batch_size=100,
                         shuffle=False, num_workers=2)

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

# 加载预训练的ResNet18模型
model = resnet18(pretrained=True)

# 修改模型以适应CIFAR-10
# 调整输入层以适应32x32的图像
model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
# 修改最后一层以适应10个类别
model.fc = nn.Linear(model.fc.in_features, 10)

# 微调策略：冻结部分层
for param in list(model.parameters())[:-10]:  # 解冻最后10层
    param.requires_grad = False

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), 
                      lr=0.001, momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)

# 设置TensorBoard写入器
timestamp = time.strftime('%Y%m%d_%H%M%S')
log_dir = os.path.join('runs', f'resnet18_cifar10_finetune_{timestamp}')
writer = SummaryWriter(log_dir)

# 训练函数
def train(epoch):
    model.train()
    train_loss = 0
    correct = 0
    total = 0
    
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_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:
            step = epoch * len(trainloader) + batch_idx
            writer.add_scalar('Train/Loss', loss.item(), step)
            writer.add_scalar('Train/Accuracy', 100.*correct/total, step)
            print(f'Epoch: {epoch} | Batch: {batch_idx+1}/{len(trainloader)} | Loss: {loss.item():.3f} | Acc: {100.*correct/total:.3f}%')
    
    # 记录每个epoch的平均训练损失和准确率
    avg_loss = train_loss / len(trainloader)
    avg_acc = 100. * correct / total
    writer.add_scalar('Epoch/Train_Loss', avg_loss, epoch)
    writer.add_scalar('Epoch/Train_Accuracy', avg_acc, epoch)
    return avg_loss, avg_acc

# 测试函数
def test(epoch):
    model.eval()
    test_loss = 0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(testloader):
            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()
    
    # 记录测试损失和准确率
    avg_loss = test_loss / len(testloader)
    avg_acc = 100. * correct / total
    writer.add_scalar('Epoch/Test_Loss', avg_loss, epoch)
    writer.add_scalar('Epoch/Test_Accuracy', avg_acc, epoch)
    
    # 记录学习率
    writer.add_scalar('Epoch/Learning_Rate', optimizer.param_groups[0]['lr'], epoch)
    
    print(f'Test Epoch: {epoch} | Loss: {avg_loss:.3f} | Acc: {avg_acc:.3f}%')
    
    # 保存最佳模型
    global best_acc
    if avg_acc > best_acc:
        print('Saving best model...')
        state = {
            'model': model.state_dict(),
            'acc': avg_acc,
            'epoch': epoch,
        }
        if not os.path.isdir('checkpoint'):
            os.mkdir('checkpoint')
        torch.save(state, './checkpoint/ckpt.pth')
        best_acc = avg_acc
    
    # 记录错误分类的图像
    if epoch % 10 == 0:  # 每10个epoch记录一次
        misclassified = []
        with torch.no_grad():
            for inputs, targets in testloader:
                inputs, targets = inputs.to(device), targets.to(device)
                outputs = model(inputs)
                _, predicted = outputs.max(1)
                mis_indices = (predicted != targets).nonzero(as_tuple=True)[0]
                
                for idx in mis_indices:
                    if len(misclassified) < 10:  # 只记录10张错误分类的图像
                        misclassified.append({
                            'image': inputs[idx],
                            'predicted': predicted[idx].item(),
                            'actual': targets[idx].item()
                        })
        
        if misclassified:
            for i, item in enumerate(misclassified):
                img = item['image'].cpu()
                writer.add_image(
                    f'Misclassified/{classes[item["actual"]]}_as_{classes[item["predicted"]]}',
                    img, epoch
                )
    
    return avg_loss, avg_acc

# 可视化模型结构
def visualize_model():
    sample_input = torch.rand(1, 3, 32, 32).to(device)
    writer.add_graph(model, sample_input)

# 可视化样本数据
def visualize_samples():
    # 获取一个批次的训练数据
    images, labels = next(iter(trainloader))
    # 创建图像网格
    img_grid = torchvision.utils.make_grid(images)
    # 添加图像网格到TensorBoard
    writer.add_image('CIFAR10_Samples', img_grid)
    # 添加标签到TensorBoard
    class_labels = [classes[i] for i in labels]
    writer.add_embedding(
        images.view(images.size(0), -1),
        metadata=class_labels,
        label_img=images,
        global_step=0
    )

# 可视化特征图
def visualize_feature_maps(inputs, epoch):
    model.eval()
    # 获取第一层卷积的特征图
    first_conv = model.conv1
    feature_maps = first_conv(inputs.to(device))
    
    # 创建特征图网格
    grid = torchvision.utils.make_grid(
        feature_maps[:8].unsqueeze(1),  # 只显示前8个特征图
        nrow=4, normalize=True, scale_each=True
    )
    writer.add_image('FeatureMaps/Conv1', grid, epoch)

# 主训练循环
if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    
    best_acc = 0
    start_epoch = 0
    
    # 可视化模型和样本
    visualize_model()
    visualize_samples()
    
    # 训练模型
    total_epochs = 50
    for epoch in range(start_epoch, total_epochs):
        print(f'Epoch {epoch}/{total_epochs-1}')
        print('-' * 10)
        
        train_loss, train_acc = train(epoch)
        test_loss, test_acc = test(epoch)
        
        # 学习率调度
        scheduler.step()
        
        # 可视化训练过程中的特征图
        if epoch % 5 == 0:  # 每5个epoch可视化一次
            inputs, _ = next(iter(testloader))
            visualize_feature_maps(inputs[:1], epoch)
        
        # 记录训练过程中的权重和梯度直方图
        if epoch % 10 == 0:  # 每10个epoch记录一次
            for name, param in model.named_parameters():
                if param.requires_grad:
                    writer.add_histogram(f'Params/{name}', param, epoch)
                    if param.grad is not None:
                        writer.add_histogram(f'Grads/{name}', param.grad, epoch)
        
        print()
    
    # 关闭TensorBoard写入器
    writer.close()
    print('Training completed!')
    print(f'Best accuracy: {best_acc:.2f}%')