【学习笔记】【Pytorch】十七、模型测试套路

一、内容概述

利用已经训练好的模型，然后给它提供输入，判断输出是否正确，即模型的应用测试。

在模型测试也会有一些坑：

• 神经网络的输入一般是4阶张量，而图片是3阶张量。
• 采用GPU训练的模型，不能直接在CPU上使用。

二、模型测试套路

在模型测试也会有一些坑：

• 神经网络的输入一般是4阶张量(batch_size, Channel, Hight, weight)，而图片是3阶张量(Channel：3, Hight：32, Weight：32)。
  解决方案：torch.reshape() API变换shape。

• 采用GPU训练的模型，不能直接在CPU上使用。

  • 解决方案1：将GPU训练的模型映射到CPU上测试。
  • 解决方案2：GPU训练的模型不用映射到CPU上，将图片映射到GPU上测试。

代码实现

import torch
import torchvision
from PIL import Image
from CIFAR_model import Model  # 导入CIFAR_model.py里的Model类定义

image_path = "./dog2.png"
image = Image.open(image_path)  # 加载3通道的图片数据，3阶张量
print(image)  # 

# 将图片 Resize(缩放) 到32x32尺寸，适合模型输入，最后在转化为Tensor实例
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
                                            torchvision.transforms.ToTensor()])
image = transform(image)
# 转化为4阶张量（模型网络的输入要求张量的阶数为4）
image = torch.reshape(image, (1, 3, 32, 32))
print(image.shape)  # torch.Size([1, 3, 32, 32])

# -----------------1.测试方式a（常用，CPU上测试）-----------------
print("\nCPU上测试:")
# 加载训练好的模型
# 采用GPU训练的模型，要想在CPU上测试，必须映射到CPU上（或者模型不用映射到CPU上，而图片映射到GPU上）
model = torch.load("CIFAR_model_9.pth", map_location=torch.device("cpu"))
# print(model)

model.eval()   # 模型进入测试模式（仅针对有Dropout,BatchNorm层的网络结构）
with torch.no_grad():  # 有利于节约内存和性能
    output = model(image)
print(output)
print(output.argmax(1))  # 方向1最大值的索引值




# -----------------2.测试方式b（GPU上测试）-----------------
print("\nGPU上测试:")
# 加载训练好的模型
model = torch.load("CIFAR_model_9.pth")
# print(model)

model.eval()   # 模型进入测试模式（仅针对有Dropout,BatchNorm层的网络结构）
with torch.no_grad():  # 有利于节约内存和性能
    # 将图片放到GPU上运行（采用GPU训练的模型，图片应该放到GPU上）
    output = model(image.cuda())
    # output = model(image.to("cuda:0"))  # 等效于上一句
print(output)
print(output.argmax(1))  # 方向1最大值的索引值

输出：

torch.Size([1, 3, 32, 32])

CPU上测试:
tensor([[-3.5857, -2.3755,  1.2455,  3.7991,  1.2300,  7.1283,  1.3590,  1.7726,
         -6.5740, -0.3369]])
tensor([5])

GPU上测试:
tensor([[-3.5860, -2.3755,  1.2455,  3.7993,  1.2299,  7.1285,  1.3591,  1.7727,
         -6.5740, -0.3368]], device='cuda:0')
tensor([5], device='cuda:0')