如何使用PyTorch的ResNet18模型对列表中的多张图像执行批量预测

报错原因 #

• 第一个报错：transforms 预处理组件仅支持单张 PIL 图像/ndarray 数组输入，无法直接处理图像列表
• 第二个报错：直接将列表转成 tensor 后维度顺序为 [batch_size, height, width, channels]，而 PyTorch CNN 要求输入维度为 [batch_size, channels, height, width]，所以模型把高度维度的36误识别为通道数，触发维度不匹配错误

实现代码 #

import torch
from torchvision import transforms

# 预定义预处理流程，无需重复创建
preprocess = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )])

# 读取图像列表，x中每个元素为(36, 60, 3)格式
x = dataset(source_p)

# 逐个处理单张图像，收集处理后的tensor
processed_imgs = []
for im in x:
    # 单张图像处理后维度为 [3, 36, 60]
    img_tensor = preprocess(im)
    processed_imgs.append(img_tensor)

# 堆叠为批量tensor，维度为 [N, 3, 36, 60]，N为列表中图像总数
batch_tensor = torch.stack(processed_imgs, dim=0)

resnet18.eval()
# 推理阶段关闭梯度计算，降低显存占用
with torch.no_grad():
    out = resnet18(batch_tensor)
# 输出out维度为[N, 1000]，对应N张图的分类结果，需要拉平可使用out.flatten(start_dim=1)

注：PyTorch 0.4及之后版本已经将Tensor和Variable接口合并，无需再使用Variable包装tensor。

内容的提问来源于stack exchange，提问作者Rayane