当前位置：首页 > article >正文

Python加载 TorchScript 格式的 ResNet18 模型分类该模型进行预测并输出预测的类别和置信度

article 2024/9/24 15:28:33

首先加载预训练的 ResNet18 模型。
将模型设置为评估模式，以确保特定层（如 Dropout 和 BatchNorm）在评估时具有确定性的行为。
创建一个形状为 (1, 3, 224, 224) 的随机张量作为示例输入。
使用 torch.jit.trace 函数追踪模型在给定示例输入上的行为，将模型转换为 TorchScript 格式。
保存 TorchScript 格式的模型为 resnet18_torchscript.pt 文件，并打印转换成功的消息。

import torch
import torchvision.models as models

# 加载预训练的 ResNet18 模型
model = models.resnet18(pretrained=True)

# 将模型设置为评估模式
model.eval()

# 创建示例输入张量
example_input = torch.rand(1, 3, 224, 224)

# 使用 torch.jit.trace 追踪模型
traced_model = torch.jit.trace(model, example_input)

# 保存 TorchScript 模型
traced_model.save('resnet18_torchscript.pt')

print("ResNet18 模型已成功转换为 TorchScript 格式并保存。")

定义图像处理函数 process_img：

- process_img 函数接受一个图像路径作为参数。
- 使用 cv2.imread 读取图像，将图像从 BGR 颜色空间转换为 RGB 颜色空间（因为很多深度学习模型期望输入为 RGB 格式）。
- 将图像的像素值归一化到 [0, 1] 范围。
- 使用 cv2.resize 将图像调整为 (224, 224) 的尺寸，这通常是 ResNet18 模型期望的输入尺寸。
- 使用 np.transpose 将图像的维度顺序从 HWC（Height-Width-Channel）转换为 CWH（Channel-Height-Width），以符合 PyTorch 的输入要求。
- 使用 np.expand_dims 在批量维度上扩展图像，使其形状变为 (1, C, H, W)。
- 最后将处理后的图像转换为 PyTorch 张量，并指定数据类型为 torch.float32，然后返回该张量。

def process_img(img_path):
    img=cv2.imread(img_path)

    img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    img = img / 255
    img=cv2.resize(img,dsize=(224,224))
    img=np.transpose(img,(2,0,1))#HWC-->CWH
    img=np.expand_dims(img,axis=0)
    img=torch.tensor(img,dtype=torch.float32)
    return img

图像预测部分：

定义一个图像路径 img_path，并将其传入 process_img 函数，得到处理后的图像张量 img。
- 使用 torch.jit.load 加载之前保存的 TorchScript 格式的模型。
- 将处理后的图像张量传入模型进行前向传播，得到输出张量 output。
- 使用 torch.argmax 在输出张量的维度 1 上找到具有最大值的索引，即预测的类别。
- 最后打印出预测的类别和对应类别的置信度（输出张量中对应类别的值）。

img_path='dog.jpg'
img=process_img(img_path)
model=torch.jit.load('resnet18_torchscript.pt')#还是torchscript格式的
output=model.forward(img)
cls=torch.argmax(output,axis=1)
print('预测的类别是：',cls.item(),'置信度是',output[0][cls].item())

预测图片