1. 项目概述与核心需求拆解1.1 为什么单独给“预测”写一篇很多做迁移学习的同学训练完模型就以为万事大吉等真正要拿模型去用的时候才发现还有一堆细节没处理。比如训练时数据是归一化过的预测时忘了对输入做同样的归一化结果准确率直接崩掉又比如训练时用的是datasets.ImageFolder自动处理的类别映射预测时手写了个类别列表顺序对不上结果长得像猫的图像被标成狗——这种问题几乎每个用ResNet做过迁移学习的人都踩过。这个系列前面几篇讲的是数据准备、网络搭建、训练与验证到了这一篇核心就是把训练好的权重加载回来写一个predict.py让它能够对单张或多张图像做推理输出类别名称和置信度。这篇的内容不复杂但每一点都可能直接决定模型落地时好不好用。predict.py全称就是“网络预测脚本”它解决的问题很明确给定一张图像加载预训练好的ResNet模型参数经过与训练时完全一致的数据预处理前向推理一次给出top-k的预测结果并把结果可视化或者写入日志方便后续集成到更大的系统里。适合谁来看如果你是刚刚跑通ResNet迁移学习训练流程、正准备把模型用起来的同学这篇文章可以直接照着抄如果你是想搞清楚推理阶段和训练阶段之间那些“隐式约定”均值方差、通道顺序、类别索引、device位置这篇文章也能帮你把坑填平。1.2 这个项目要用到的技术栈脚本虽然不大涉及的知识点却很集中。我简单列一下后面会逐一展开PyTorch的模型定义与权重加载包括torch.load、load_state_dict以及state_dict的键名匹配问题。数据预处理Pipeline的复现也就是transforms.Normalize的均值和标准差必须和训练时保持一致。类别索引与真实标签的映射关系包括从ImageFolder导出的class_to_idx怎么用。GPU与CPU环境的兼容处理包括模型和输入张量要放在同一个设备上。推理时的模型模式切换model.eval()和torch.no_grad()的意义以及为什么不能省。后处理与可视化包括softmax的数值稳定性、top-k的取法以及把预测结果画到图像上的常规做法。这些点单独拎出来都不难但组合在一起时先后顺序和细节处理就非常重要。下面我会按照一个标准predict.py从零到能用的顺序把每个环节讲透。2. 整体设计思路与方案选型2.1 为什么选择“脚本化”而不是“Jupyter Notebook”方式我见过不少同学在Notebook里做预测把一段推理代码反复执行改来改去。这在实验阶段没问题但真要到了部署或者批量处理图像的阶段脚本化是更合理的选择。原因很简单脚本可以接收命令行参数可以批量处理文件可以嵌入到更大的流程里而且没有Notebook那些“隐性状态”的干扰——你不需要关心“上一个Cell是不是跑过了”只要从零开始执行结果就是确定的。predict.py的设计目标就三个输入一张图像路径输出预测结果支持批量预测结果清晰可读。这三个目标决定了脚本的整体骨架参数解析、模型加载、数据预处理、推理循环、结果输出。2.2 权重文件的管理方式训练完模型之后通常会有多个权重文件比如best_model.pth、last_model.pth。predict.py必须允许用户指定要加载哪一个权重文件所以我推荐把权重路径做成命令行参数--weights而不是硬编码在脚本里面。这样你在跑不同实验对比时只需要切换参数不需要改代码。关于权重文件的格式我个人更推荐在训练时直接保存整个模型torch.save(model.state_dict(), path)而不是保存整个model对象。前者只保存参数文件更小加载时对模型类定义的依赖更小后者如果你改了模型类的定义老权重可能就废了。predict.py这边对应使用load_state_dict配合你自己重新构建的模型结构这样最稳妥。2.3 类别映射的持久化难题训练时如果你用了torchvision.datasets.ImageFolder加载数据它会自动根据子文件夹的名称生成类别索引顺序是字典序。比如你的训练集目录长这样data/ train/ cat/ dog/那么class_to_idx就会是{cat: 0, dog: 1}。这个映射关系在训练时在dataset.classes和dataset.class_to_idx里都能拿到但到了预测阶段你手上通常只有一张图像不太可能为了一张图再构建一个完整的ImageFolder数据集。所以我的建议是在训练脚本里把class_to_idx保存成一个JSON文件比如class_indices.json和权重文件放在同一个目录下。predict.py读取这个JSON建立一个“索引-类别名”的映射预测时把模型输出的类别索引转换成人能看懂的字符串。千万别自己手写一个列表一旦顺序写反所有预测结果全是错的而且这种错误非常隐蔽单看准确率数字根本发现不了。2.4 预处理的一致性没有商量余地迁移学习的预训练权重比如在ImageNet上训练过的ResNet官方给定的预处理方法是Resize(256)、CenterCrop(224)、ToTensor()、Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225])。训练时你的数据增强流水线里可能还有RandomResizedCrop、RandomHorizontalFlip但预测阶段绝对不能用随机增强必须用确定性的缩放和裁剪否则同一张图每次预测出来的结果可能都不一样这对生产环境是致命的。还有一个容易被忽略的点如果你的训练脚本里为了某种原因修改了归一化的均值或标准差那么预测脚本必须同步修改。不要看到网上某个代码用的是[0.5, 0.5, 0.5]就直接复制你得确认它跟你的模型训练时的配置是同一套。这就像你用同一把钥匙去开两把不同的锁锁芯不对钥匙再高级也没用。3. 核心细节解析与实操要点3.1 模型构建与权重加载predict.py里构建模型的方式必须和训练时的模型结构完全一致。以ResNet18为例假设你的分类任务是10类训练时最后一层全连接是model.fc nn.Linear(512, 10)那么预测脚本里也要做同样的替换。如果你在训练时还加了Dropout或者其他自定义层这些全部要复现。权重加载的标准写法是model models.resnet18(pretrainedFalse, num_classes10) model.load_state_dict(torch.load(best_model.pth, map_locationcpu))这里有两个细节值得说清楚。第一pretrainedFalse。推理阶段不需要再下载ImageNet预训练权重因为我们要加载的是自己训练好的权重。如果写成pretrainedTrue脚本会先下载一份ImageNet权重然后再被你自己的权重覆盖白白浪费时间而且在没网的环境下直接报错。第二map_locationcpu。这个参数非常关键。如果训练时用的是GPU保存的权重张量会带有cuda:0的设备信息预测时如果机器上没有可用的GPU直接torch.load会报错。加上map_locationcpu之后权重会被加载到CPU内存里后续你再决定是把模型放到GPU还是留在CPU。加载权重之后强烈建议做一步校验state_dict torch.load(best_model.pth, map_locationcpu) model.load_state_dict(state_dict)如果训练时保存的是model.state_dict()那键名会和模型的参数名完全匹配。但如果你保存的是model本身加载时就得用torch.load直接拿到model对象这种情况下load_state_dict就不好使了。我的建议是从训练阶段开始就统一用state_dict让这条链路简单直接。3.2 设备管理不该省的判断脚本里最好有一个自动判断设备的逻辑device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)这样同一个脚本在有无GPU的机器上都能跑。有一个常见的低级错误是模型在GPU上输入图像张量在CPU上前向传播时直接报错“Expected all tensors to be on the same device”。反过来也一样。所以图像张量也要做个.to(device)。另外如果是在GPU上做推理别忘了每个batch处理完之后及时释放显存。虽然单张图像推理通常不会爆显存但如果你批量处理成千上万张图中间变量会不断累积。用with torch.no_grad():包裹推理部分能减少很大一部分显存占用因为不再保存梯度的中间张量。3.3 数据预处理Pipeline的复现预测脚本里的预处理流程transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这个Pipeline看起来简单每一步都有它存在的理由。Resize(256)先把图像缩放到短边256CenterCrop(224)再从中心裁剪出224x224的区域。224是ResNet标准输入尺寸ImageNet预训练时用的就是这个分辨率。如果你直接Resize(224)而不做CenterCrop实际上是把图像非等比压缩了物体比例会变形干扰模型判断。这个细节在ImageNet权重上可能只影响一两个百分点的精度但既然标准流程已经给出照做就好。ToTensor()做的事情是把PIL图像转成Tensor并且把像素值从0~255缩放到0.0~1.0。这一步容易误解总觉得“归一化”已经包括了缩放其实ToTensor()和Normalize()是两个独立的操作。前者负责数值范围后者负责按通道减去均值除以标准差。这里再提醒一次预测时的transforms必须是确定性的不能包含任何随机操作。如果你在训练时用了RandomResizedCrop(224, scale(0.08, 1.0))做增强预测时也照抄了这个那就完全错了。预测时图像每次都会被随机裁掉一块输出结果可能每次都不一样而且这种错误很不好排查因为你复现不了固定的输出。3.4 处理PIL图像与OpenCV图像的通道差异predict.py里读取图像一般用PILfrom PIL import Image img Image.open(image_path).convert(RGB)convert(RGB)一定要加因为如果图像是RGBA格式或者灰度图Image.open返回的模式可能不同后续ToTensor()虽然能处理但通道数不一致会引发模型输入维度错误。加上convert(RGB)强制统一成三通道干净利落。如果你习惯用OpenCV的cv2.imread请注意OpenCV读进来的是BGR顺序的numpy数组直接拿来预测会得到错误结果。你需要先cv2.cvtColor(img, cv2.COLOR_BGR2RGB)然后Image.fromarray(rgb_img)再走后面的Pipeline。这是个非常经典的坑几乎每隔一段时间就会有人在社区问一遍。3.5 输入张量的形状问题模型输入是四维张量(batch_size, channels, height, width)。单张图像做完ToTensor()之后是三维的(3, 224, 224)所以必须增加一个batch维度input_tensor transform(img).unsqueeze(0).to(device)很多第一次写推理代码的人会漏掉unsqueeze(0)这一步结果模型前向传播直接报维度错误。这个错误信息可能不那么直观提示的是“Expected 4D input, got 3D input”只要看到这个就该反射性地想到batch维度丢了。4. 实操过程与关键环节实现4.1 完整的predict.py代码框架直接给一个能跑的框架你可以在此基础上按自己的实际情况调整。import argparse import json import torch import torch.nn as nn from torchvision import models, transforms from PIL import Image def load_class_indices(json_path): with open(json_path, r) as f: class_to_idx json.load(f) idx_to_class {int(v): k for k, v in class_to_idx.items()} return idx_to_class def build_model(num_classes, weights_path, device): model models.resnet18(pretrainedFalse, num_classesnum_classes) state_dict torch.load(weights_path, map_locationcpu) model.load_state_dict(state_dict) model model.to(device) model.eval() return model def preprocess_image(image_path): transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(image_path).convert(RGB) return transform(img).unsqueeze(0) def predict(model, image_path, idx_to_class, device, top_k5): input_tensor preprocess_image(image_path).to(device) with torch.no_grad(): outputs model(input_tensor) probs torch.softmax(outputs, dim1)[0] top_probs, top_indices torch.topk(probs, top_k) results [] for prob, idx in zip(top_probs, top_indices): class_name idx_to_class[idx.item()] results.append((class_name, prob.item())) return results def main(): parser argparse.ArgumentParser(descriptionResNet 迁移学习预测脚本) parser.add_argument(--image, typestr, requiredTrue, help输入图像路径) parser.add_argument(--weights, typestr, requiredTrue, help模型权重文件路径) parser.add_argument(--class_indices, typestr, requiredTrue, help类别映射JSON文件路径) parser.add_argument(--topk, typeint, default5, help输出前K个预测结果) args parser.parse_args() device torch.device(cuda if torch.cuda.is_available() else cpu) idx_to_class load_class_indices(args.class_indices) model build_model(len(idx_to_class), args.weights, device) results predict(model, args.image, idx_to_class, device, top_kargs.topk) print(f预测图像: {args.image}) for i, (class_name, prob) in enumerate(results): print(fTop-{i1}: {class_name}, 置信度: {prob:.4f}) if __name__ __main__: main()这个脚本的结构是参数解析 - 加载类别映射 - 构建模型并加载权重 - 逐张图像前向推理 - 打印top-k结果。命令行的用法python predict.py --image ./test/cat.001.jpg --weights ./checkpoints/best_model.pth --class_indices ./checkpoints/class_indices.json --topk 34.2 类别映射JSON的生成方式class_indices.json从哪来我建议在训练脚本里训练结束后把dataset.class_to_idx保存下来import json class_to_idx train_dataset.class_to_idx with open(class_indices.json, w) as f: json.dump(class_to_idx, f)这样生成的JSON内容大概是{cat: 0, dog: 1, bird: 2}注意JSON里的键是字符串值是整数。读取时直接json.load会得到同样的字典但整数类型在反复读写的过程中可能变成坑所以我上面的代码里才专门做了int(v)转换确保索引是整型。4.3 批量预测与结果保存上面给的是单张图像版本实际使用中更多时候是要处理一整个文件夹的图像。这时候不建议在命令行里一张一张调用脚本效率太低。更合理的做法是让predict.py也支持--image_dir参数对目录下的所有图像做循环推理结果统一写到CSV文件里。批量版本的核心逻辑import os import csv from PIL import Image def predict_folder(model, image_dir, idx_to_class, device, top_k1): image_paths [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.lower().endswith((.jpg, .jpeg, .png))] with open(predict_results.csv, w, newline) as csvfile: writer csv.writer(csvfile) writer.writerow([image_path, predicted_class, confidence]) for path in image_paths: results predict(model, path, idx_to_class, device, top_ktop_k) class_name, prob results[0] writer.writerow([path, class_name, f{prob:.4f}]) print(f{path} - {class_name} ({prob:.4f}))一个提醒批量推理时如果没有任何打印信息跑起来会让人非常焦虑不知道是卡住了还是在正常推进。所以无论如何都要加个进度提示。也可以考虑引入进度条库但简单场景下自己打印足够用。4.4 可视化把预测结果画在图像上只看打印出来的文字还是不够直观尤其是当你要向别人展示模型效果时最好能把类别名和置信度直接画到图像上。一个简单的做法from PIL import ImageDraw, ImageFont def visualize(image_path, results): img Image.open(image_path).convert(RGB) draw ImageDraw.Draw(img) class_name, prob results[0] text f{class_name}: {prob:.2f} draw.text((10, 10), text, fill(255, 0, 0)) img.save(predicted.jpg)这里有个小细节ImageDraw.Draw.text()默认使用PIL内置字体字号很小而且在不支持中文的环境下中文字符直接显示为方块。如果你想在图像上显示中文类别名需要加载一个中文字体文件font ImageFont.truetype(/path/to/NotoSansCJK-Regular.ttc, size30) draw.text((10, 10), text, fontfont, fill(255, 0, 0))我实际使用时的经验是先把类别名和置信度在图上渲染好再保存为带_pred后缀的新文件不要覆盖原图。这样原图还保留着方便对比和复查。5. 常见问题与排查技巧实录5.1 问题速查表我整理了平时被问到最多的问题按出现的频率排了个序每个都给出了现象、原因和解决办法。问题现象可能原因解决方法Expected all tensors to be on the same device模型在GPU输入在CPU对输入张量调用.to(device)size mismatch for fc.weight模型全连接输出维度与权重不匹配检查num_classes与训练时保持一致预测结果全部是同一个类别类别映射顺序搞反了检查class_indices.json是否与训练时一致预测准确率远低于训练集验证准确率预处理步骤与训练不一致检查Resize/Crop/Normalize参数是否照搬训练增强RuntimeError: CUDA out of memory批量处理时显存累积使用torch.no_grad()或减小batch中文类别显示为方块PIL默认字体不支持中文加载中文字体文件IndexError: list index out of range类别索引超出映射范围确认idx_to_class字典的键覆盖所有类别5.2 权重加载报错Missing key(s)与Unexpected key(s)这是所有迁移学习新手最容易碰到的报错。核心原因是模型结构和权重文件里保存的key集合对不上。如果你训练时用了DataParallel包装模型保存的state_dict键名会多出module.前缀比如module.fc.weight而你自己构建的模型里键名是fc.weight。解决方式有两种第一种在预测脚本里加载后做一次字符串处理from collections import OrderedDict new_state_dict OrderedDict() for k, v in state_dict.items(): name k[7:] if k.startswith(module.) else k new_state_dict[name] v model.load_state_dict(new_state_dict)第二种更推荐训练时保存之前先取原始模型的state_dicttorch.save(model.module.state_dict() if hasattr(model, module) else model.state_dict(), best_model.pth)这样保存下来的键名就是干净的预测脚本不需要额外处理。类似的还有类别数量不匹配问题。训练时是1000类预测时num_classes10而权重文件的fc.weight形状是(1000, 512)load_state_dict直接报size mismatch。你需要确认自己模型构建的num_classes与训练时完全一致。5.3 softmax到底是加上还是不加很多人在预测时会把模型的原始输出直接打印发现数值可能是正的也可能是负的看着很奇怪。记住模型最后的全连接层输出是logits不是概率。要想得到0到1之间的概率值必须做softmax。PyTorch的torch.softmax(outputs, dim1)实现已经考虑了数值稳定性你只需要指定正确的dim。dim1指的是类别那一维因为outputs的形状是(batch_size, num_classes)。如果你写成dim0softmax会沿着batch维度算结果每个样本的概率之和不为1完全错误。如果对数值稳定性特别在意可以用torch.log_softmax加torch.exp的组合但对推理来说直接softmax已经足够稳定性能影响可以忽略。5.4 模型推理速度太慢怎么办如果你对单张图像的预测耗时特别敏感有几个可以立竿见影的优化手段。首先确定模型在推理模式下运行。model.eval()不仅关掉Dropout和BatchNorm的训练行为而且配合torch.no_grad()能显著减少计算和内存开销。很多人忘了model.eval()模型里如果有BatchNorm层推理时的统计量会被batch数据影响导致结果不稳定。其次如果你的环境支持GPU但输入图像很小、批次也只有1张GPU的优势可能发挥不出来因为数据从CPU拷到GPU的时间抵消了计算加速。这种情况下适当增大batch比如一次同时喂8张或16张图像吞吐量会明显提升。最后如果确实需要极致的单张推理速度可以考虑把模型导出为TorchScript或者ONNX格式用专门推理引擎加速。但这是另一个话题了predict.py这个阶段先把功能做对做稳性能优化可以后面再谈。5.5 数据流中的“隐藏约定”我把这类问题统称为“隐藏约定”因为它们在代码里不会显式出现但一旦破坏结果就全错。通道顺序约定RGB还是BGR模型训练时用的是哪种预测时就必须是哪种。像素范围约定0~1还是0~255Normalize之前的数值范围必须匹配。图像方向约定要不要做水平翻转训练时如果用了RandomHorizontalFlip模型见过翻转后的图但预测时通常不做翻转这是允许的因为模型学到的是翻转不变的特征。缩放方式约定是等比缩放还是直接拉伸ResNet官方流程用的是等比缩放加中心裁剪如果你直接拉伸到224长宽比失真对识别效果有影响。常见的坑是训练代码里Normalize用的mean和std是自己在数据集上统计出来的比如[0.5, 0.5, 0.5]但预测脚本直接从网上复制了ImageNet的[0.485, 0.456, 0.406]。两个模型用不同的预处理方式训练出来预测时用错的那套结果自然不会好。我的习惯是把预处理参数写成全局常量放在训练脚本和预测脚本都要导入的一个公共配置文件中。这样只维护一份两边的值不会漂移。如果你已经写好了两个独立脚本至少要做到把预处理代码从训练脚本里原样复制到预测脚本中不要手打不要凭记忆。6. 实操心得与扩展建议6.1 我在调试predict.py时踩过的坑写这个脚本最让我印象深刻的一次调试是某次模型预测结果里猫的图像被识别成狗而且置信度非常高。我一度怀疑是模型学坏了后来仔细排查才发现训练时数据集的类别顺序是{dog: 0, cat: 1}而我在写预测脚本时“按常理”认为0是猫1是狗把映射列表写反了。模型输出的索引完全正确只是我把索引翻译成了错误的类别名。那之后我给自己定了一条规矩任何预测脚本里类别索引映射必须从训练阶段的class_to_idx导入绝不手写。如果历史上没有保存class_indices.json那就重新加载一次训练数据集从train_dataset.class_to_idx现场导出。这样虽然多了一步但可以保证万无一失。还有一次同事在预测一批灰度图像时发现部分图像预测出来的类别全是同一个。排查到最后发现是因为部分图像在保存时用了错误的通道顺序某些图像实际是BGR顺序但后缀名是.jpgPIL的convert(RGB)并不会重新排列通道它只是把已有的三个通道映射过去。这个问题的本质是数据源本身不干净脚本里能做的最好事前检查是打印一下图像的模式和通道均值确认它确实是正常的RGB图像。6.2 predict.py如何接入更大的服务predict.py只是个起点。实际生产环境中你可能要把模型封装成一个HTTP接口或者嵌入到实时视频流处理程序里。这时候有两个地方值得提前设计。一个是模型的初始化方式。不要在每次请求时都重新加载权重这非常耗时。更好的做法是进程启动时加载一次模型后续请求直接复用。用PyTorch写服务时这通常意味着在全局作用域初始化模型然后在处理函数里只做前向推理。另一个是输入图像的尺寸动态性。如果你的系统可能收到不同尺寸的图像预处理Pipeline里的Resize(256)就是固定的没问题但如果你希望支持多尺度输入就要在服务层面做适配。对于ResNet这类固定输入尺寸的模型最省事的方式是统一在预处理阶段缩放到224x224不要试图去改模型结构。后期的方向还可以做模型量化、模型蒸馏甚至把多个模型串成一条推理流水线。但这些都是后话先把predict.py写对、写稳后面每一步才有坚实的地基。6.3 给初学者的最后提醒如果你正在照着写自己的predict.py我希望你记住三点第一训练脚本和预测脚本里的预处理、类别映射、模型结构这三件事必须保持严格一致。不一致的地方就是预测结果崩掉的地方。第二不要迷信网上复制来的代码。理解每一行在做什么尤其是transforms.Normalize里的均值和标准差是从哪个数据集来的、为什么是这三个数这些都必须搞清楚。第三推理代码要尽量简单、鲁棒、可复现。宁可多写几行判断也不要让脚本在边角情况下崩溃。比如图像路径不存在、权重文件加载失败、类别映射为空这些都要有明确的错误提示信息而不是让Python抛出一堆堆栈然后退出。predict.py虽然是一个很小很具体的脚本但它是模型从训练走向实用的最后一公里。把这一公里走稳后面无论是做部署、做服务、做产品都会顺手很多。