Day40 预训练
预训练 = 用别人在超大数据集（ImageNet）上训练好的模型权重 → 拿来给我们用，少训练、效果好、速度快
import torch import torch.nn as nn import torch.optim as optim from torchvision import models, datasets, transforms from torch.utils.data import DataLoader import matplotlib.pyplot as plt # ===================== 1. 配置 ===================== device = torch.device("cuda" if torch.cuda.is_available() else "cpu") batch_size = 64 lr = 0.001 epochs = 5 num_classes = 10 # ===================== 2. 数据预处理（必须匹配预训练模型要求） ===================== transform = transforms.Compose([ transforms.Resize((224, 224)), # ResNet 标准输入 224 transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], # ImageNet 均值 std=[0.229, 0.224, 0.225]) ]) # CIFAR10 数据集 train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) test_dataset = datasets.CIFAR10(root='./data', train=False, transform=transform) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) # ===================== 3. 加载预训练模型（核心！） ===================== # 加载 ResNet18 预训练权重 model = models.resnet18(pretrained=True) # ===================== 4. 冻结卷积层（只训练最后的全连接层） ===================== for param in model.parameters(): param.requires_grad = False # 冻结所有层 # 替换最后一层全连接层（1000类 → 10类） in_features = model.fc.in_features model.fc = nn.Linear(in_features, num_classes) model = model.to(device) # ===================== 5. 损失函数 & 优化器 ===================== criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.fc.parameters(), lr=lr) # 只训练最后一层 # ===================== 6. 训练函数 ===================== def train(model, loader, criterion, optimizer, device): model.train() total_loss = 0 correct = 0 for data, target in loader: data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() total_loss += loss.item() pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() avg_loss = total_loss / len(loader) acc = 100.0 * correct / len(loader.dataset) return avg_loss, acc # ===================== 7. 测试函数 ===================== def test(model, loader, criterion, device): model.eval() total_loss = 0 correct = 0 with torch.no_grad(): for data, target in loader: data, target = data.to(device), target.to(device) output = model(data) total_loss += criterion(output, target).item() pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() avg_loss = total_loss / len(loader) acc = 100.0 * correct / len(loader.dataset) return avg_loss, acc # ===================== 8. 开始训练 ===================== print("===== 开始训练（仅训练最后一层）=====") for epoch in range(1, epochs+1): train_loss, train_acc = train(model, train_loader, criterion, optimizer, device) test_loss, test_acc = test(model, test_loader, criterion, device) print(f"Epoch {epoch:2d} | " f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | " f"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}%") # ===================== 9. 解冻全部层，进行微调（进阶） ===================== print("\n===== 解冻全部层，进行精细微调 =====") for param in model.parameters(): param.requires_grad = True # 解冻所有层 optimizer = optim.Adam(model.parameters(), lr=1e-5) # 小学习率 for epoch in range(1, 3): train_loss, train_acc = train(model, train_loader, criterion, optimizer, device) test_loss, test_acc = test(model, test_loader, criterion, device) print(f"Epoch {epoch:2d} | " f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | " f"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}%")@浙大疏锦行