ConvNeXt 的 torchvision 版本 推理实现

ConvNeXt 的 torchvision 版本 推理实现flyfishimporttorchimportosfromPILimportImagefromtorchvisionimporttransforms# 引入你自定义的ConvNeXt模型文件确保和本文件同目录fromconvnext_tinyimportconvnext_tiny# 配置 # 本地模型权重路径.pth / .pt / .pth.tar 格式MODEL_WEIGHT_PATHconvnext_tiny-983f1562.pth# 测试图像路径TEST_IMAGE_PATHtest.jpg# 分类类别数必须和训练模型时的num_classes一致NUM_CLASSES1000# 输出Top-K预测结果TOPK5# # 设备配置自动使用GPU无GPU则用CPUDEVICEtorch.device(cudaiftorch.cuda.is_available()elsecpu)defget_convnext_preprocess(): ConvNeXt 官方标准图像预处理必须严格匹配训练流程 步骤Resize → CenterCrop → ToTensor → Normalize preprocesstransforms.Compose([# 缩放到256像素transforms.Resize(256),# 中心裁剪为224x224ConvNeXt-Tiny标准输入尺寸transforms.CenterCrop(224),# 转为Tensor数值归一化到 [0, 1]transforms.ToTensor(),# ImageNet 标准归一化ConvNeXt训练使用的均值/方差transforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])])returnpreprocessdefload_local_model(model_path:str,num_classes:int): 初始化自定义ConvNeXt模型 加载本地权重 :param model_path: 本地权重文件路径 :param num_classes: 分类类别数 :return: 评估模式的模型 # 1. 初始化模型结构modelconvnext_tiny(num_classesnum_classes)# 2. 检查权重文件是否存在ifnotos.path.exists(model_path):raiseFileNotFoundError(f模型权重不存在路径{model_path})# 3. 加载本地权重map_location适配CPU/GPUcheckpointtorch.load(model_path,map_locationDEVICE)# strictFalse适配修改了分类头num_classes的模型避免权重不匹配报错model.load_state_dict(checkpoint,strictFalse)# 4. 模型迁移到设备 开启评估模式关闭dropout/bn训练特性modelmodel.to(DEVICE)model.eval()returnmodeldefimage_preprocess(image_path:str,transform): 图像完整预处理加载 → 转RGB → 预处理 → 增加Batch维度 → 设备迁移 # 检查图像是否存在ifnotos.path.exists(image_path):raiseFileNotFoundError(f测试图像不存在路径{image_path})# 1. 加载图像 强制转为RGB解决灰度图/透明通道报错问题imageImage.open(image_path).convert(RGB)# 2. 应用预处理 [H, W, C] → [C, H, W]tensor_imagetransform(image)# 3. 增加Batch维度[C, H, W] → [1, C, H, W]模型要求批量输入tensor_imagetensor_image.unsqueeze(0)# 4. 迁移到设备tensor_imagetensor_image.to(DEVICE)returnimage,tensor_imagedefmodel_infer(model,input_tensor): 模型推理 后处理 1. 无梯度推理节省显存/加速 2. Softmax将Logits转为概率 3. 提取Top-K置信度和类别索引 withtorch.no_grad():# 禁用梯度计算推理必备# 前向传播输出原始预测值 (logits)outputsmodel(input_tensor)# 后处理1softmax归一化得到0~1的概率值probabilitiestorch.softmax(outputs,dim1)# 后处理2获取Top-K 类别索引 置信度topk_probs,topk_indicestorch.topk(probabilities,kTOPK)# 转为numpy格式方便打印topk_probstopk_probs.cpu().numpy()[0]topk_indicestopk_indices.cpu().numpy()[0]returntopk_probs,topk_indicesdefprint_result(topk_probs,topk_indices,class_namesNone):打印最终预测结果print(\n*50)print(f图像分类预测结果Top-{TOPK})print(*50)ifclass_namesisNone:class_names[f类别_{i}foriinrange(NUM_CLASSES)]fori,(idx,prob)inenumerate(zip(topk_indices,topk_probs)):print(fTop{i1}{class_names[idx]:10}| 置信度{prob:.4f}({prob*100:.2f}%))print(*50)if__name____main__:# 1. 初始化预处理工具 transformget_convnext_preprocess()print(图像预处理配置完成)# 2. 加载本地模型 print(f正在加载本地模型{MODEL_WEIGHT_PATH})modelload_local_model(MODEL_WEIGHT_PATH,NUM_CLASSES)print(模型加载完成已切换到评估模式)# 3. 图像预处理 print(f正在处理图像{TEST_IMAGE_PATH})raw_image,input_tensorimage_preprocess(TEST_IMAGE_PATH,transform)print(图像预处理完成)# 4. 模型推理 print(开始模型推理...)topk_probs,topk_indicesmodel_infer(model,input_tensor)print(推理完成)# 5. 打印结果 # 替换为真实类别名称# 示例custom_classes [cat, dog, car, tree]custom_classesNoneprint_result(topk_probs,topk_indices,custom_classes)输出图像预处理配置完成 正在加载本地模型convnext_tiny-983f1562.pth 模型加载完成已切换到评估模式 正在处理图像test.jpg 图像预处理完成 开始模型推理... 推理完成图像分类预测结果Top-5Top1类别_209|置信度0.3588(35.88%)Top2类别_227|置信度0.2372(23.72%)Top3类别_168|置信度0.0283(2.83%)Top4类别_208|置信度0.0222(2.22%)Top5类别_166|置信度0.0149(1.49%)初始化准备 → 模型加载 → 图像预处理 → 前向推理 → 后处理输出整体推理总流程从__main__入口开始代码会按顺序执行 5 个步骤初始化图像预处理规则加载本地权重构建并初始化 ConvNeXt 模型读取测试图像完成标准化预处理输入模型执行前向推理得到原始预测值对预测结果做后处理格式化打印 Top-K 分类结果分阶段说明阶段1全局配置与预处理初始化对应代码顶部配置项 get_convnext_preprocess()这是推理的前置准备目标是保证推理时的图像处理逻辑和模型训练时完全一致否则会出现精度大幅下降。固定全局超参数指定权重路径、测试图路径、分类类别数、输出 Top-K 数量自动选择运行设备有 GPU 则用 CUDA 加速没有则自动回退到 CPU保证代码可跨环境运行。定义标准预处理流水线按顺序执行 4 步操作和 ImageNet 训练流程严格对齐Resize(256)把图像短边缩放到 256 像素保持宽高比不变CenterCrop(224)从缩放后的图像中心裁剪出 224×224 的正方形匹配 ConvNeXt-Tiny 的标准输入尺寸ToTensor()把 PIL 图像0~255 整数HWC 格式转为 PyTorch 张量0~1 浮点数CHW 格式Normalize用 ImageNet 数据集的均值和标准差做通道归一化把数据分布对齐到训练时的分布是保证精度的关键一步。阶段2模型加载与初始化对应函数load_local_model()目标是构建和权重匹配的网络结构载入预训练参数切换到推理专用状态。构建模型骨架调用convnext_tiny(num_classes1000)按照我们之前分析的 CNBlock 四阶段结构实例化一个随机初始化的 ConvNeXt-Tiny 网络。载入本地权重文件用torch.load读取.pth权重文件map_locationDEVICE保证权重自动加载到对应设备避免 GPU 权重在 CPU 上报错model.load_state_dict(checkpoint, strictFalse)把权重参数填入模型骨架。strictFalse是容错设计如果修改了分类头类别数、或有少量层不匹配不会直接报错只加载匹配的层适配微调后的模型推理。切换到推理模式model.to(DEVICE)把模型整体迁移到 GPU/CPUmodel.eval()开启评估模式这是推理的必要操作。它会关闭所有训练专用的随机组件比如StochasticDepth随机深度、Dropout 都会停止工作保证每次推理结果稳定、可复现同时 BatchNorm/LayerNorm 会用训练好的全局统计量而非 batch 统计量。阶段3图像预处理与格式对齐对应函数image_preprocess()目标是把一张普通图片转换成模型能接收的标准输入张量。图像读取与格式统一Image.open(...).convert(RGB)强制转为 3 通道 RGB 格式兼容灰度图、带透明通道的 RGBA 图避免通道数不匹配报错。应用预处理流水线调用阶段1定义的 transform完成缩放、裁剪、转张量、归一化输出形状为[3, 224, 224]的张量通道 × 高 × 宽。补充 Batch 维度unsqueeze(0)把形状从[C, H, W]扩展为[1, C, H, W]。PyTorch 模型默认接受批量输入第一维是 batch 大小单张图推理也要补一个 batch1 的维度否则维度不匹配会报错。数据迁移到对应设备把图像张量移到和模型相同的设备GPU/CPU保证运算时设备一致。阶段4模型前向推理对应函数model_infer()这是计算环节执行前向传播得到预测结果同时做推理优化。禁用梯度计算with torch.no_grad():是推理的标准写法推理不需要反向传播关闭梯度计算可以大幅节省显存、提升推理速度避免推理过程中意外修改模型参数。前向传播得到原始输出outputs model(input_tensor)调用模型的forward方法数据依次经过 Stem → 4个Stage → 全局平均池化 → 分类头最终输出形状为[1, 1000]的原始预测值logits数值没有归一化不代表概率。后处理转概率 取Top-Ktorch.softmax(outputs, dim1)对 1000 个类别的原始输出做归一化得到 0~1 之间的概率值所有类别概率和为 1torch.topk(probabilities, k5)从 1000 个类别中选出置信度最高的 5 个返回对应的概率值和类别索引。结果转numpy格式把张量从 GPU 移回 CPU转为 numpy 数组方便后续打印输出。阶段5结果格式化输出对应函数print_result()把 Top-K 的索引和置信度转换成可读的结果打印出来如果传入了自定义类别名称列表就显示具体类别名默认显示类别_索引的占位符同时打印置信度的小数和百分比形式直观展示预测可信度。