PyTorch微调实战：用ResNet18快速搞定你的第一个图像分类任务（附完整代码）

PyTorch微调实战用ResNet18快速搞定你的第一个图像分类任务当你第一次面对一个图像分类任务时预训练模型就像是一位经验丰富的导师能帮你快速入门。ResNet18作为经典的卷积神经网络在保持轻量级的同时提供了出色的特征提取能力。本文将带你从零开始完成一个完整的图像分类项目。1. 环境准备与数据加载在开始之前确保你已经安装了PyTorch和torchvision。如果你使用conda环境可以通过以下命令安装conda install pytorch torchvision torchaudio -c pytorch对于图像分类任务数据准备是第一步。假设我们要处理一个花卉分类数据集包含5个类别。PyTorch提供了方便的ImageFolder加载器可以自动处理按文件夹分类的图像数据。from torchvision import datasets, transforms # 定义数据增强和标准化 data_transforms { train: transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), val: transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } # 加载数据集 data_dir flower_data image_datasets {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in [train, val]} dataloaders {x: torch.utils.data.DataLoader(image_datasets[x], batch_size4, shuffleTrue, num_workers4) for x in [train, val]} dataset_sizes {x: len(image_datasets[x]) for x in [train, val]} class_names image_datasets[train].classes提示数据增强是提升模型泛化能力的关键但验证集不需要随机增强只需进行中心裁剪和标准化。2. 模型加载与微调策略ResNet18在ImageNet上预训练的特征提取器已经学会了识别各种视觉模式。我们只需要针对特定任务调整最后的全连接层。import torchvision.models as models # 加载预训练模型 model models.resnet18(pretrainedTrue) # 冻结所有卷积层参数 for param in model.parameters(): param.requires_grad False # 替换最后的全连接层 num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, len(class_names)) # 将模型移动到GPU device torch.device(cuda:0 if torch.cuda.is_available() else cpu) model model.to(device)这里我们采用了局部微调策略只训练最后的全连接层。这种方法的优势在于训练速度快因为大部分参数不需要更新需要的训练数据量较少避免了在小数据集上过度拟合如果你有足够的数据和计算资源也可以尝试分层微调给不同层设置不同的学习率# 分层微调示例 ignored_params list(map(id, model.fc.parameters())) base_params filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer optim.SGD([ {params: base_params, lr: 1e-4}, {params: model.fc.parameters(), lr: 1e-2} ], momentum0.9)3. 训练过程与技巧训练神经网络需要选择合适的损失函数和优化器。对于分类任务交叉熵损失是标准选择。criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.fc.parameters(), lr0.001, momentum0.9) # 学习率调度器 exp_lr_scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size7, gamma0.1)完整的训练循环包含以下几个关键步骤def train_model(model, criterion, optimizer, scheduler, num_epochs25): for epoch in range(num_epochs): print(fEpoch {epoch}/{num_epochs - 1}) print(- * 10) # 每个epoch都有训练和验证阶段 for phase in [train, val]: if phase train: model.train() # 设置训练模式 else: model.eval() # 设置评估模式 running_loss 0.0 running_corrects 0 # 迭代数据 for inputs, labels in dataloaders[phase]: inputs inputs.to(device) labels labels.to(device) # 梯度清零 optimizer.zero_grad() # 前向传播 with torch.set_grad_enabled(phase train): outputs model(inputs) _, preds torch.max(outputs, 1) loss criterion(outputs, labels) # 只在训练阶段反向传播和优化 if phase train: loss.backward() optimizer.step() # 统计 running_loss loss.item() * inputs.size(0) running_corrects torch.sum(preds labels.data) if phase train: scheduler.step() epoch_loss running_loss / dataset_sizes[phase] epoch_acc running_corrects.double() / dataset_sizes[phase] print(f{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}) return model注意在验证阶段要设置model.eval()这会关闭dropout和batch normalization的随机性确保评估结果的一致性。训练过程中常见的几个问题及解决方案损失不下降检查学习率是否合适确认数据加载是否正确检查模型是否真的在更新参数过拟合增加数据增强添加正则化如权重衰减减少模型复杂度训练不稳定使用学习率预热尝试不同的优化器如Adam梯度裁剪4. 模型评估与可视化训练完成后我们需要评估模型在测试集上的表现。除了准确率混淆矩阵能提供更详细的分类情况。from sklearn.metrics import confusion_matrix import seaborn as sn import pandas as pd # 在验证集上评估 model.eval() all_preds [] all_labels [] with torch.no_grad(): for inputs, labels in dataloaders[val]: inputs inputs.to(device) labels labels.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) # 构建混淆矩阵 cm confusion_matrix(all_labels, all_preds) df_cm pd.DataFrame(cm, indexclass_names, columnsclass_names) plt.figure(figsize(10,7)) sn.heatmap(df_cm, annotTrue, fmtd) plt.xlabel(Predicted) plt.ylabel(True) plt.show()可视化模型关注的特征区域也很有帮助可以使用类激活映射CAM技术from torchcam.methods import CAM # 创建CAM提取器 cam_extractor CAM(model, layer4) # 获取单张图像 for inputs, _ in dataloaders[val]: input_tensor inputs[0].unsqueeze(0).to(device) break # 前向传播 out model(input_tensor) activation_map cam_extractor(out.squeeze(0).argmax().item(), out) # 可视化 plt.imshow(input_tensor.cpu().squeeze().permute(1,2,0)) plt.imshow(activation_map[0].squeeze().cpu().numpy(), alpha0.5, cmapjet) plt.show()5. 模型保存与部署训练好的模型可以保存下来供后续使用。PyTorch提供了几种保存方式# 保存整个模型 torch.save(model, flower_classifier.pth) # 只保存模型参数推荐 torch.save(model.state_dict(), flower_classifier_params.pth) # 加载模型 model models.resnet18(pretrainedFalse) model.fc nn.Linear(model.fc.in_features, len(class_names)) model.load_state_dict(torch.load(flower_classifier_params.pth)) model model.to(device)对于生产环境部署可以考虑以下方案TorchScript将模型转换为脚本形式可以在没有Python环境的情况下运行ONNX开放神经网络交换格式支持跨框架部署Flask/Django构建Web API服务一个简单的Flask部署示例from flask import Flask, request, jsonify from PIL import Image import io app Flask(__name__) model.eval() app.route(/predict, methods[POST]) def predict(): if file not in request.files: return jsonify({error: no file uploaded}) file request.files[file].read() image Image.open(io.BytesIO(file)) # 预处理 input_tensor data_transforms[val](image).unsqueeze(0).to(device) # 预测 with torch.no_grad(): output model(input_tensor) _, pred torch.max(output, 1) return jsonify({class: class_names[pred.item()]}) if __name__ __main__: app.run(host0.0.0.0, port5000)6. 进阶技巧与优化当基本模型运行起来后你可以尝试以下方法进一步提升性能学习率搜索def find_lr(model, criterion, optimizer, dataloader, init_value1e-8, final_value10.0): lr init_value optimizer.param_groups[0][lr] lr best_loss float(inf) losses [] lrs [] for inputs, labels in dataloader: inputs inputs.to(device) labels labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) # 提前停止如果损失爆炸 if loss 4 * best_loss: break if loss best_loss: best_loss loss losses.append(loss.item()) lrs.append(lr) loss.backward() optimizer.step() lr * (final_value / init_value) ** (1/len(dataloader)) optimizer.param_groups[0][lr] lr return lrs, losses混合精度训练from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for inputs, labels in dataloader: inputs inputs.to(device) labels labels.to(device) optimizer.zero_grad() with autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()模型剪枝from torch.nn.utils import prune # 随机剪枝 parameters_to_prune ( (model.conv1, weight), (model.layer1[0].conv1, weight), ) prune.global_unstructured( parameters_to_prune, pruning_methodprune.L1Unstructured, amount0.2, )在实际项目中我发现数据质量往往比模型结构更重要。花费时间清洗和增强数据集通常比调整超参数带来的提升更明显。另外使用wandb或TensorBoard记录实验过程能帮助你更好地理解模型行为。