Pytorch 完整的模型验证套路

Pytorch 完整的模型验证套路

利用已经测试好的模型，然后给它提供输入

我们从网上随便找一张小狗的图片，放在 imgs 文件夹下

新建一个 model_test 文件

from PIL import Image

import torch
import torchvision.transforms
from torch import nn

image_path = "imgs/dog.png"
image = Image.open(image_path)          # 打开文件
# image = image.convert('RGB')          # png 格式是四个通道，除了RGB三通道外，还有一个透明通道。所以我们调用 convert('RGB')，保留其颜色通道

transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),    # Compose() 将几个变换联立在一起
                                            torchvision.transforms.ToTensor()])
image = transform(image)

# 创建网络模型
class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4, 64),
            nn.Linear(64, 10)
        )

    def forward(self, input):
        input = self.model(input)
        return input

# 加载网络模型, 为什么要写前面的类？
# 因为 model_0.pth 用第一种保存方法, 有缺陷
model = torch.load("model_0.pth", map_location=torch.device('cpu'))     # gpu 上训练的模型要映射到cpu上 ，或者用 .cuda() 也可以
print(model)

image = torch.reshape(image, (1, 3, 32, 32))

model.eval()                                    # 将模型转换为测试模型
with torch.no_grad():                           # 可以节约我们的内存，提高性能
    #   output = model(image.cuda())            # gpu 上训练的模型也要在gpu上加载
    output = model(image)
print(output)

print(output.argmax(1))

会出现如下结果：

tensor([6]) 表示是第六类，即 frog-青蛙，预测失败了，艹

再测试一下：

将图片换成飞机：

image_path = "imgs/airplane.png"

再运行一下啊：

tensor([0]) 表示预测是第 0 件，即飞机，预测成功 !