【Pytorch】用自己训练的resnet18模型进行推理

在网上看了这篇博客：

【Pytorch】使用ResNet-50迁移学习进行图像分类训练
https://blog.csdn.net/heiheiya/article/details/103028543

有的小伙伴就有疑问，模型训练好后，怎么进行推理。于是，我写了这篇关于使用自己训练的resnet18模型进行推理，训练的部分请参考上面那篇。

# *_* coding : UTF-8 *_*
# 开发人员： csu·pan-_-||
# 开发时间： 2020/12/29 19:16
# 文件名称： resnet_battery_infer.py
# 开发工具： PyCharm
# 功能描述： 用自己训练好的resnet18模型进行推理

import torch
import time
import numpy as np
import os
import cv2

modelPath =r'Battery\battery_resnet18.pt'  # 模型路径
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
validPath = r'Battery\valid\ok' # 需要测试的图片路径
txtPath = 'Battery/results'     # 结果存储路径
files = os.listdir(validPath)   # 展开图片文件列表
model = torch.load(modelPath)   # 加载模型
class_name = ['NG','OK']        # 类别名称
classList = []    # 方便计数

with torch.no_grad():
    model.eval()
    start = time.time()
    for j, input in enumerate(files):
        onestart = time.time()
        img = cv2.imread(os.path.join(validPath,input))
        # 转换图片格式
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x800x800
        img = np.ascontiguousarray(img)
        # 使用cuda
        img = torch.from_numpy(img).to(device)
        img = img.float()  # int转换为float
        img /= 255.0  # 归一化
        if img.ndimension() == 3:
            img = img.unsqueeze(0)   # 1x3x800x800
        outputs = model(img)
        ret, predictions = torch.max(outputs.data, 1)
        print('outputs.data: ',outputs.data)
        # outputs.data:  tensor([[6.4012e-04, 9.9936e-01]], device='cuda:0')
        # tensor转换成list:
        inferclass = predictions.cpu().numpy().tolist()[0]
        classList.append(class_name[inferclass])
        print('class: {:s}'.format(class_name[inferclass]))
        print('input: ',input)
        # class: 0 或 1
        oneend = time.time()
        print('one last: {:.4f}'.format(oneend-onestart))
        with open(txtPath + '/' + str(j) + '.txt', 'w') as f:
            f.write(class_name[inferclass])   # 存储结果为 'NG' 或 'OK'

    end = time.time()
    print('all last: {:.4f}s in {:d} imgs'.format(end - start,len(files)))
    # 统计类别数量：
    print('NG num: ',classList.count('NG'))
    print('OK num: ', classList.count('OK'))

在我的显卡上：TITAN RTX，还是蛮快的，推理一张图大概 14ms

上面那个推理是针对归一化到[0,1]模型的，如果归一化到了[-1,1]，则推理时也要对应归一化到[-1,1]

norm_mean = [0.485, 0.456, 0.406]    # 均值
norm_std = [0.229, 0.224, 0.225]     # 方差

inference_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std),   # 转换到-1 - 1
])

def img_transform(img_rgb, transform=None):
    """
    将数据转换为模型读取的形式
    :param img_rgb: PIL Image
    :param transform: torchvision.transform
    :return: tensor
    """
    if transform is None:
        raise ValueError("找不到transform！必须有transform对img进行处理")
    img_t = transform(img_rgb)
    return img_t

用以下方法读图和推理，这里需要用到一个新库：PIL：

# path --> img
    img_rgb = Image.open(os.path.join(validPath,input)).convert('RGB')
    # img --> tensor
    img_tensor = img_transform(img_rgb, inference_transform)
    img_tensor.unsqueeze_(0)    # 1x3x800x800
    img_tensor = img_tensor.to(device)   # use cuda
    outputs = model(img_tensor)  # 推理