预训练 用别人在超大数据集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, trainTrue, downloadTrue, transformtransform) test_dataset datasets.CIFAR10(root./data, trainFalse, transformtransform) train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue) test_loader DataLoader(test_dataset, batch_sizebatch_size, shuffleFalse) # 3. 加载预训练模型核心 # 加载 ResNet18 预训练权重 model models.resnet18(pretrainedTrue) # 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(), lrlr) # 只训练最后一层 # 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(dim1) 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(dim1) 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, epochs1): train_loss, train_acc train(model, train_loader, criterion, optimizer, device) test_loss, test_acc test(model, test_loader, criterion, device) print(fEpoch {epoch:2d} | fTrain Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | fTest 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(), lr1e-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(fEpoch {epoch:2d} | fTrain Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | fTest Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}%)浙大疏锦行