117_PyTorch 实战:利用训练好的模型进行单张图片验证

训练模型(Train)和测试模型(Test)最终都是为了应用(Inference) 。本篇将教你如何加载已经保存的 .pth 模型文件,并用一张外部图片来检验它的分类能力。

1. 验证流程的三大核心步骤

  1. 准备测试环境:包括加载模型、处理单张图片。
  2. 图像预处理 :输入的图片必须经过与训练时完全相同的缩放(Resize)和归一化(ToTensor)。
  3. 模型推理:将图片送入模型,并解析输出结果。

2. 代码实战:验证模型的分类结果

文件演示了如何加载一个在 CIFAR-10 上训练好的模型,并识别一张"狗"或"飞机"的图片。

Python

复制代码
import torch
import torchvision
from PIL import Image
from torch import nn

# 1. 读取外部图片
image_path = "dog.png" # 或者是你本地的图片路径
img = Image.open(image_path)
# 如果是 RGBA 格式(带透明度),需转为 RGB
img = img.convert('RGB')

# 2. 图像预处理(必须与训练时保持一致)
transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32, 32)),
    torchvision.transforms.ToTensor()
])
img = transform(img)
# 增加 Batch 维度,从 [3, 32, 32] 变为 [1, 3, 32, 32]
img = torch.reshape(img, (1, 3, 32, 32))

# 3. 加载训练好的模型
# 注意:如果模型是用方式一保存的,加载时需要能访问到网络定义类
model = torch.load("tudui_29.pth", map_location=torch.device('cpu'))

# 4. 进入验证模式
model.eval()
with torch.no_grad():
    output = model(img)

# 5. 解析输出
# output 是一个长度为 10 的向量,值最大的位置即为预测类别
print(output)
predict_idx = output.argmax(1).item()

# CIFAR-10 的类别映射(固定顺序)
classes = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
print(f"预测结果是:{classes[predict_idx]}")

3. 关键细节分析

为什么需要 img.convert('RGB')?

PNG 图片通常包含 4 个通道(RGBA),其中 A 是透明度。而我们的模型是基于 RGB 3 通道训练的。使用 convert('RGB') 可以保证无论输入什么格式的图片,都能适配模型。

map_location 的妙用

如果你的模型是在 GPU 上训练保存的,但现在你想在只有 CPU 的电脑上运行验证,加载时必须加上 map_location=torch.device('cpu'),否则会报错。

argmax(1) 的逻辑

模型输出的是图片属于这 10 个类别的"概率得分"。通过 argmax(1),我们能直接提取出得分最高的那个位置的索引(例如 5 代表狗)。


4. 总结:从模型到应用

分析完这个文件,我们就完成了从数据采集到实战部署的全流程:

  1. 训练集/测试集->训练出高准确率模型。
  2. 保存模型-> 持久化存储。
  3. 单张验证->将模型应用到真实场景中。

💡 学习小结

当你能成功地输入一张从网上下载的图片,并让模型正确报出"cat"或"dog"时,你就真正完成了深度学习的闭环。

相关推荐
懂点Tech1 小时前
ctfshow misc28:GIF 高度修改解题与脚本实现
开发语言·windows·python
comedate1 小时前
SMLib 的 Examples 导读:PyAPI / SM_API / SMLib 三组示例怎么读
python·cad·smlib·examples·sm_api
云和数据.ChenGuang2 小时前
Python 重写(Override)深度解析:从 MRO 到 super() 的完整机制
java·开发语言·人工智能·python·langchain·1024程序员节
西西木科技丨Shopify开发机构2 小时前
CalDigit TS4案例:AI产品图的Shopify整站接入与验收
人工智能·shopify·独立站
川石课堂软件测试2 小时前
学习软件测试需掌握的一些知识
数据库·python·sql·功能测试·网络协议·mysql·http
2601_965448872 小时前
AI视频制作:做系列节气民俗视频音色不统一,AI 音色克隆省事的方法
人工智能·机器学习·音视频
美狐美颜SDK开放平台2 小时前
视频美颜SDK试用DEMO阶段应该测试什么?直播APP接入实操指南
开发语言·人工智能·深度学习·音视频·美颜sdk·直播美颜sdk·第三方美颜sdk
HRaitest3 小时前
秋招技术岗招聘提效:三类 AI 招聘方案的技术底层与选型逻辑
大数据·人工智能·求职招聘·ai-native
天远数科3 小时前
零信任架构实战:基于天远企业股权变更构建自动化商业银行对公授信存续期股权异动与资产抽逃预警网关
运维·人工智能·架构·自动化
万化冥合3 小时前
PyBox屏幕键盘如何辅助演示键盘操作?
python·计算机外设