1. 为什么是Resnet18而不是Resnet50或VGG16——模型选型与整体方案设计1.1 残差结构到底解决了什么问题做图像多分类最常见的一个困惑是网络越深效果越好那直接上Resnet50甚至Resnet101不就行了我在自己的数据集上对比过多次实际结论和很多人想的不太一样。Resnet18属于Resnet家族里最轻量的一员核心是残差学习。它的基础模块是BasicBlock由两个3x3卷积组成通过一个跳跃连接shortcut把输入直接加到输出上。这个设计解决的是网络退化问题——网络加深之后训练误差反而上升这不是过拟合而是优化难度变大。残差结构让每一层学习的不再是完整的映射而是输入和输出之间的差值也就是残差。这样一来梯度可以沿着shortcut直接回传哪怕网络再深梯度也不会轻易消失。对比VGG16VGG是纯卷积堆叠没有残差连接训练时对学习率和初始化非常敏感收敛速度也明显慢。我最初用VGG16训过一版同样的数据集和epochsVGG16大概需要多花1.5到2倍的时间才能达到相近的精度而且显存占用高出一截。那为什么不直接上Resnet50对个人学习者或者中小型项目来说Resnet50参数量约2550万是Resnet18的将近3倍。如果你的训练集只有几千张甚至几百张图Resnet50很容易过拟合而且训练速度大幅下降。Resnet18参数量只有约1120万在CPU上也能勉强跑推理在GTX 1660这种级别的显卡上训练一张224x224的图batch_size设为32大概只需要不到1秒一个batch。对大多数应用场景来说Resnet18的精度已经足够尤其是迁移学习场景下它的泛化能力往往比在ImageNet上预训练的大模型更稳。1.2 多分类任务的完整技术链路图像多分类任务听起来高大上拆开其实就是几个固定环节数据集准备、数据预处理、模型定义、损失函数、训练循环、验证评估、模型保存、推理测试。每个环节都有对应的坑这篇文章主要针对PyTorch框架展开。PyTorch做这件事有几个天然优势动态计算图让调试非常直观torchvision库里直接集成了Resnet18的预训练权重一句话就能加载另外DataLoader的数据加载机制对自定义数据集非常友好。1.3 环境与依赖安装我的运行环境如下供参考Python 3.8PyTorch 1.102.x版本完全兼容本文代码torchvision 0.11CUDA 10.2及以上没有GPU可以跑CPU版但训练速度会慢几十倍opencv-python、numpy、matplotlib、Pillow安装核心依赖pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python numpy matplotlib pillow tqdm提示如果你用的是PyTorch 2.xtorchvision中Resnet18的权重入口略有变化建议直接使用weightsResNet18_Weights.DEFAULT参数而不是旧版的pretrainedTrue新版虽然保留兼容但会弹警告。后文代码我统一用官方推荐的新写法。关于GPU验证跑代码前务必先执行这一句import torch print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU mode)返回True才有意义否则下面所有训练代码默认走CPU跑一个epoch可能就要几十分钟很多人以为代码卡死了其实是没注意到自己在用CPU跑。2. 自己的数据集怎么准备才靠谱——目录规则、预处理与数据增强2.1 数据集目录组织最简单实用的ImageFolder方式训练自己的数据集第一步不是写代码而是把数据整理好。PyTorch的torchvision.datasets.ImageFolder要求数据集遵循一个文件夹对应一个类别的结构这是最主流也是最省事的方式。dataset/ train/ class_a/ 001.jpg 002.jpg ... class_b/ 001.jpg 002.jpg ... class_c/ ... val/ class_a/ ... class_b/ ... class_c/ ...这里有几个容易被忽略的细节第一每个类别下的图片数量尽量均衡。如果class_a有1000张class_b只有50张模型会严重偏向class_a准确率虚高但实际泛化能力很差。如果数据实在不均衡至少要在后文介绍训练参数时用WeightedRandomSampler做类别平衡或者使用带权重的损失函数。第二图片格式建议统一为.jpg或.png。混合格式不是不行只是某些损坏文件会导致读取报错排查起来非常痛苦。第三文件名任意但不要有中文和特殊字符Linux服务器上中文路径经常出问题。第四train和val的数据来源不要有重叠。很多人直接把同一批图片又分了一部分出来做验证集这会导致验证集准确率虚高因为模型已经见过了这些图。2.2 自定义Dataset类还是ImageFolderImageFolder适合数据集已经按类别分好目录的情况而且它自动根据文件夹名称生成类别索引比如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]) ]) }但真实场景往往是数据是乱序的比如一堆图片在一个文件夹类别信息在CSV里。这时候ImageFolder就不适用了必须自己写一个Dataset子类。核心就是重写三个方法__init__、__len__、__getitem__。__getitem__返回的是(image_tensor, label)的元组把图片读取和预处理都封装到这个类里。篇幅原因我在后面给出完整代码时会同时兼容这两种场景但建议新手优先使用ImageFolder代码少、不易错。2.3 数据增强为什么重要、怎么配置数据增强是提高模型泛化能力最有效的手段尤其当你只有几百张图的时候。直接让模型去学习原始图片它会记住训练集的纹理、颜色、背景一换环境就拉胯。Resnet18默认输入是224x224训练集增强我这里用了RandomResizedCrop和RandomHorizontalFlip。前者随机裁剪不同区域并缩放模拟了目标物体在画面中位置和大小变化后者是水平翻转增加了样本多样性。增强策略不要盲目堆叠。过度的增强等于给模型加噪声训练很难收敛。常见的增强手段还包括增强方式使用场景注意事项RandomResizedCrop目标物体位置不固定配合scale参数控制裁剪比例RandomHorizontalFlip左右对称的物体文字类、车牌类不要用ColorJitter对光照敏感的任务强度控制在0.2-0.3RandomRotation物体方向不确定配合fill参数填充背景Normalize所有场景必须用ImageNet均值和标准差ColorJitter是我后期给效果不佳的数据集加的一招比如transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.1)它对提升模型对光照变化、颜色偏移的鲁棒性很有帮助代价是训练会变慢一点点。2.4 数据划分比例一般建议训练集占80%验证集占10%测试集占10%。如果你没有专门留测试集也至少要把训练和验证分开。验证集的作用是监视训练过程中的过拟合决定什么时候保存模型或者停训。训练集再大验证集也不能省。3. 核心代码逐段拆解——从DataLoader到训练循环3.1 DataLoader构建与参数理解DataLoader是PyTorch的数据加载核心它帮你自动做batch划分、乱序打散、多进程并行读取。理解了它训练代码就能写得干净。from torch.utils.data import DataLoader image_datasets { train: datasets.ImageFolder(dataset/train, data_transforms[train]), val: datasets.ImageFolder(dataset/val, data_transforms[val]) } dataloaders { train: DataLoader(image_datasets[train], batch_size32, shuffleTrue, num_workers4), val: DataLoader(image_datasets[val], batch_size32, shuffleFalse, num_workers4) } batch_size {train: 32, val: 32}shuffleTrue在训练时非常重要它防止模型学到数据的固定顺序。验证集不需要shuffle因为验证不涉及梯度更新顺序不影响结果。num_workers是并行读取进程数Windows上设置过大容易报错通常设置为0或2。Linux服务器设成CPU核数减1也可以。batch_size的选择需要权衡越大GPU利用率越高、训练越稳定但显存占用也大越小梯度更新方向越不稳定但显存压力小。Resnet18在224x224分辨率下batch_size32在6GB显存显卡上基本没问题。3.2 模型构建迁移学习、冻结卷积层Resnet18的经典玩法是迁移学习——用别人在ImageNet上训练好的权重做初始化然后再用自己的数据微调。这样哪怕你只有几百张图也能快速收敛到不错的效果。完整代码import torch import torch.nn as nn import torch.optim as optim import torchvision from torchvision import models, datasets, transforms from torchvision.models import ResNet18_Weights import time import copy device torch.device(cuda:0 if torch.cuda.is_available() else cpu) print(Using device:, device) # 数据集增强与归一化 data_transforms { train: transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.1), 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]) ]) } # 加载数据集ImageFolder方式 image_datasets { train: datasets.ImageFolder(dataset/train, data_transforms[train]), val: datasets.ImageFolder(dataset/val, data_transforms[val]) } # 通过classes属性获取类别名 class_names image_datasets[train].classes num_classes len(class_names) print(类别数:, num_classes, 类别:, class_names) dataloaders { train: DataLoader(image_datasets[train], batch_size32, shuffleTrue, num_workers4), val: DataLoader(image_datasets[val], batch_size32, shuffleFalse, num_workers4) }模型定义部分这是全篇最关键的地方def build_resnet18(num_classes, feature_extractTrue): # 新写法加载预训练权重 model models.resnet18(weightsResNet18_Weights.DEFAULT) # 是否冻结卷积层参数 if feature_extract: for param in model.parameters(): param.requires_grad False # 替换最后的全连接层 in_features model.fc.in_features # Resnet18是512 model.fc nn.Linear(in_features, num_classes) return model model build_resnet18(num_classesnum_classes, feature_extractFalse) model model.to(device)feature_extract这个参数决定了是否冻结卷积层。如果训练数据很少或者你的图片和ImageNet分布差异不大建议冻结卷积层只训练最后的全连接层。这样速度快、不容易过拟合。但如果你的数据比较特殊比如是医学影像、卫星图这类就需要全网络微调效果会好很多。这个参数我做成开关大家按需调整。3.3 训练循环与验证循环实现训练循环是整篇代码的主干我把它拆成几个模块来讲方便看懂和修改。损失函数用交叉熵损失CrossEntropyLoss它是多分类的标准选择。优化器用带动量的SGD。注意一点优化器只作用于requires_grad为True的参数如果冻结了卷积层就只更新fc层的参数。criterion nn.CrossEntropyLoss() optimizer optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr0.001, momentum0.9) scheduler optim.lr_scheduler.StepLR(optimizer, step_size7, gamma0.1)学习率0.001微调是常见经验值。如果你从头训练不加载预训练权重建议学习率设置到0.01到0.1。用预训练权重时0.001比较稳太大容易把预训练学好的特征破坏掉。StepLR每7个epoch把学习率缩小10倍目的是在训练后期用更小的步长微调模型。训练与验证循环def train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochs25): since time.time() best_model_wts copy.deepcopy(model.state_dict()) best_acc 0.0 train_loss_list [] val_acc_list [] for epoch in range(num_epochs): print(fEpoch {epoch1}/{num_epochs}) print(- * 30) # 每个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) epoch_loss running_loss / len(dataloaders[phase].dataset) epoch_acc running_corrects.double() / len(dataloaders[phase].dataset) print(f{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}) if phase val: val_acc_list.append(epoch_acc.item()) if epoch_acc best_acc: best_acc epoch_acc best_model_wts copy.deepcopy(model.state_dict()) # 保存最优模型 torch.save(model.state_dict(), best_model_resnet18.pth) print(f模型已保存当前最优准确率: {best_acc:.4f}) scheduler.step() print() time_elapsed time.time() - since print(f训练完成总耗时 {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s) print(f最优验证准确率: {best_acc:.4f}) # 加载最优模型参数 model.load_state_dict(best_model_wts) return modelmodel.train()和model.eval()的区别必须搞清楚。train()模式下BatchNorm会更新均值和方差、Dropout生效eval()模式下BatchNorm用训练集统计的全局均值方差Dropout关闭。如果验证时忘了切到eval()模式验证集准确率会偏低且不稳定。optimizer.zero_grad()是每个batch开始前必须调用把上一步累加的梯度清零。忘了这行代码梯度会在batch之间累积loss会震荡甚至爆炸。torch.set_grad_enabled(phase train)是一个很巧妙的写法验证阶段不需要计算梯度这样可以节省显存、加速推理。等价写法是把验证部分包在with torch.no_grad():里。preds取的是模型输出在类别维度上的最大值索引也就是预测的类别编号。torch.max(outputs, 1)这里的1表示沿着类别维度取最大值。执行训练num_epochs 25 model train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochsnum_epochs)3.4 保存最优模型与继续训练的说明上面的代码在验证集准确率刷新时自动保存了best_model_resnet18.pth。这里有一个很多人没注意到的细节保存的是model.state_dict()还是整个模型我习惯保存state_dict因为这样最节省空间而且加载时只需先构建同样的模型结构再load。保存整个模型虽然加载方便但会把优化器状态、模型类定义都打包进去遇到PyTorch版本升级容易出兼容问题。如果训练中途断了还想接着原来的状态继续那就要把优化器状态也保存下来。你的checkpoint应该像这样checkpoint { epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict(), best_acc: best_acc } torch.save(checkpoint, checkpoint.pth)加载时对应恢复即可。不过日常使用中最优模型已经能满足绝大多数需求checkpoint一般用在超长训练任务上。4. 训练时的超参数设置与踩坑排查4.1 学习率、batch_size、epochs到底怎么定这部分是新手最容易迷茫的地方。直接给经验值参数取值范围我的推荐说明学习率0.0001 ~ 0.010.001迁移学习用0.001从头训练用0.01batch_size8 ~ 12832按显存定6G显卡32起步epochs20 ~ 5025看验证集acc是否还有上升趋势momentum0.8 ~ 0.990.9SGD标配weight_decay0 ~ 0.0010.0005加一点正则化防过拟合epochs不要固定死。正确做法是观察验证集准确率曲线如果还在上升就继续训练如果连续多个epoch不再上升甚至下降就该停了。最简单的方式是在代码里加入早停机制例如连续10个epoch验证准确率没有刷新记录就停止训练。学习率是影响训练效果最大的超参数。早年我犯过一个错误用0.1的学习率去微调Resnet18结果loss直接发散越训越差。原因是0.1的步长太大把预训练模型已经学好的特征全部冲乱了。用0.001微调收敛曲线非常平滑。如果发现loss下降缓慢可以尝试学习率预热warm-up前5个epoch从0.0001线性增加到0.001然后再正常训练。4.2 常见的三个坑过拟合、loss不降、显存不足过拟合的表现是训练集准确率很高验证集却上不去。解决思路按优先级排列第一增加数据增强强度第二减少模型容量比如从Resnet50换回Resnet18第三降低全连接层之后的过拟合风险在fc层前加Dropout第四增大weight_decay第五增加训练数据量。loss不降通常是这几个原因学习率太大导致震荡、数据预处理有问题比如忘了Normalize、类别标签错位、代码忘了切换train/eval模式。我还遇到过一种情况就是多人数据集拼接时类别索引对不上训练集和验证集同一种类的标签编号不一样结果loss永远降不下来。显存不足报错是CUDA out of memory。处理办法是降低batch_size这是最快最有效的。如果降到8还不够就降低输入分辨率到160x160。再不行就检查是否有别的进程占着显存nvidia-smi训练到一半显示OOM往往是因为验证阶段忘了包torch.no_grad()导致验证时也计算了梯度显存翻倍。4.3 用matplotlib画训练曲线训练曲线能直观反映模型状态。把训练loss和验证准确率画出来过拟合一目了然loss下降但验证集准确率停滞甚至下降就是过拟合的信号。画图的代码很简单import matplotlib.pyplot as plt def plot_curves(train_losses, val_accs): fig, ax1 plt.subplots() ax1.plot(train_losses, b-, labeltrain loss) ax1.set_xlabel(Epoch) ax1.set_ylabel(Loss, colorb) ax2 ax1.twinx() ax2.plot(val_accs, r-, labelval acc) ax2.set_ylabel(Accuracy, colorr) plt.title(Training Loss and Validation Accuracy) fig.tight_layout() plt.savefig(training_curves.png, dpi150) plt.show()我在训练代码里已经收集了train_loss_list和val_acc_list训练完成后直接调用即可。5. 真实测试与评估——别被训练准确率骗了5.1 加载模型预测单张图片训练完成之后最关键的一步是用训练时没见过的图片测试。把模型加载回来写一个预测脚本跟训练代码是完全独立的可以是单独一个Python文件。import torch import torchvision.transforms as transforms from PIL import Image from torchvision import models from torchvision.models import ResNet18_Weights def load_model(num_classes, model_pathbest_model_resnet18.pth): model models.resnet18(weightsResNet18_Weights.DEFAULT) in_features model.fc.in_features model.fc torch.nn.Linear(in_features, num_classes) model.load_state_dict(torch.load(model_path, map_locationtorch.device(cpu))) model.eval() return model def predict_image(model, image_path, class_names): # 必须和训练时的预处理一致 transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) image Image.open(image_path).convert(RGB) image_tensor transform(image).unsqueeze(0) with torch.no_grad(): outputs model(image_tensor) _, predicted torch.max(outputs, 1) probs torch.softmax(outputs, dim1) class_idx predicted.item() confidence probs[0][class_idx].item() return class_names[class_idx], confidence if __name__ __main__: class_names [class_a, class_b, class_c] # 改成你的类别名和训练时打印的一致 model load_model(num_classeslen(class_names)) label, conf predict_image(model, test_images/sample1.jpg, class_names) print(f预测结果: {label}, 置信度: {conf:.2f})这里有两个隐藏的坑。第一model.load_state_dict加载前必须先把模型结构构建出来并且fc层的输出维度要和保存时的一致。如果你训练时是3分类这里写5分类会直接报错key不匹配。第二加载模型后必须调用model.eval()。我自己就栽过这个跟头模型加载完直接预测结果每次输出都不一样一开始以为是随机性后来才想起忘了切eval模式BatchNorm还在用当前batch的统计量导致推理结果不稳定。第三图片预处理必须和训练时验证集一致。很多人训练时用CenterCrop(224)预测时却直接resize到224导致图片分辨率分布改变了精度会掉几个点。5.2 不只是准确率——多分类要会看混淆矩阵多分类任务光看准确率远远不够。如果你的数据有10个类别准确率90%你根本不知道是哪个类别容易混淆。混淆矩阵是排查分类错误的最好工具。from sklearn.metrics import confusion_matrix, classification_report import numpy as np def evaluate_model(model, dataloader, device, class_names): model.eval() all_preds [] all_labels [] with torch.no_grad(): for inputs, labels in dataloader: inputs inputs.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, target_namesclass_names)) cm confusion_matrix(all_labels, all_preds) print(混淆矩阵:) print(cm) return cmclassification_report里的精确率precision和召回率recall是发现问题的关键。如果某个类别recall很低说明很多该类的图片被误判成了别的类别这时候需要检查是不是该类训练样本太少了或者该类别和其他类别在视觉上太相似。混淆矩阵能直接告诉你模型到底在哪些类别之间摇摆这是针对性优化最重要的依据。6. 最后再分享几个实战中的细节6.1 数据量极少怎么办——微调策略调整如果每个类别只有几十张图训练数据严重不足第一时间考虑用数据增强解决第二时间考虑冻结卷积层只训练fc层。我把feature_extract设为True时通常只需要训练10个epoch就能达到不错的效果。训练速度也快因为卷积层的参数不参与反向传播显存占用直接砍半。如果数据还是不够可以试试让模型在大数据集上先做无监督预训练或者使用更强的数据增强如CutMix、MixUp但这些方法复杂度较高建议先把前面的基础方案跑通。6.2 迁移学习的边界情况——图像风格差异大时有一种容易翻车的情况你的图像和ImageNet训练集的图像风格差异很大。比如ImageNet的图大多是自然的物体照片而你要分类的是X光片、图纸、手绘图案这类直接用预训练权重不一定占优势。我测试过一批纹理差异特别大的数据迁移学习的效果居然不如随机初始化从头训因为预训练特征跟目标数据完全不匹配起不到正向作用。判断方法是对比两种初始化的验证集准确率。如果预训练权重表现明显更差就干脆从零训练学习率设为0.01epochs适当增加。如果不知道自己的数据算不算差异大就先跑预训练看结果再决定。6.3 推理速度优化思路模型训好之后如果要做部署Resnet18的推理速度在CPU上大约几十毫秒到一百多毫秒一张图基本可以满足简单实时需求。如果还嫌慢可以试试把输入尺寸从224降到160精度损失大约1-2个点速度提升近一倍。再往后就是模型量化和剪枝那就是另一个话题了。实际跑一遍下来你会发现Resnet18并不是最花哨的网络却是把准确率、速度、显存占用、上手难度各方面权衡得最均衡的选择。用它训练自己的数据集既能学到图像分类的完整流程又不至于被复杂模型拖住脚步后面想换Resnet50或者Resnet101只需要替换一行模型定义代码训练流程完全不用动。希望这篇文章能帮你少走一些弯路顺利跑通自己的第一个图像多分类模型。