一学就会的深度学习基础指令及操作步骤（4）模型部署推理

模型部署

模型加载

将之前训练好的模型从文件加载到内存中，并准备好用于预测或继续训练

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = torch.load('model.pth', map_location=device)

加载处理预测数据

# 输入须与训练模型所用数据的形状相匹配
IMG_WIDTH = 28  
IMG_HEIGHT = 28

# 经过预处理后，将相同大小和颜色的图像传入模型
preprocess_trans = transforms.Compose([
    transforms.ToDtype(torch.float32, scale=True), # 将图像转换为浮点型，将 `scale` 设置为 `True` 以将 [0, 255] 转换为 [0, 1]
    transforms.Resize((IMG_WIDTH, IMG_HEIGHT)),    # 将图像调整为 28 x 28 像素
    transforms.Grayscale()  # 将图像转换为灰度（训练图片是灰度图）
])

查看图片尺寸

processed_image = preprocess_trans(image)
processed_image.shape
>>> torch.Size([1, 28, 28])

模型期望接收一批图像，.unsqueeze(0)，第一个维度通常是批次维度
squeeze 移除 1 个维度
unsqueeze 在我们指定的索引处添加 1 个维度

batched_image = processed_image.unsqueeze(0)
batched_image.shape
>>> torch.Size([1, 1, 28, 28])

预测结果

output = model(batched_image_gpu)
prediction = output.argmax(dim=1).item()
alphabet = "abcdefghiklmnopqrstuvwxy"
alphabet[prediction]

流程总结

def predict_letter(file_path):
    # Show image
    show_image(file_path)
    # Load and grayscale image
    image = tv_io.read_image(file_path, tv_io.ImageReadMode.GRAY)
    # Transform image
    image = preprocess_trans(image)
    # Batch image
    image = image.unsqueeze(0)
    # Send image to correct device
    image = image.to(device)
    # Make prediction
    output = model(image)
    # Find max index
    prediction = output.argmax(dim=1).item()
    # convert prediction to letter
    predicted_letter = alphabet[prediction]
    return predicted_letter```