Python训练营打卡Day45

作业：对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 = 0correct = 0total = 0for 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_idxwriter.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 / totalwriter.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 = 0correct = 0total = 0with 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 / totalwriter.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_accif 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)# 添加图像网格到TensorBoardwriter.add_image('CIFAR10_Samples', img_grid)# 添加标签到TensorBoardclass_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.conv1feature_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 = 0start_epoch = 0# 可视化模型和样本visualize_model()visualize_samples()# 训练模型total_epochs = 50for 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}%')