Pytorch学习笔记--模型的测试/验证

最新推荐文章于 2025-08-07 18:05:35 发布

原创最新推荐文章于 2025-08-07 18:05:35 发布 · 514 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#学习 #笔记

要明确，test.py 就是把模型应用到实际场景中。test.py在py文件夹下

先把要识别的图片放在文件夹images中。目录结构如下：

输入图片，调整图片格式以便于输入到模型，（此模型输入为3通道的32*32），再加载之前训练好的模型（训练十轮的aa_9.pth）：

#test就是把模型应用到实际场景中

import torch
import torch.nn as nn
import torchvision
from PIL import Image



image_path="../images/dog.png"
image=Image.open(image_path)
print(image)
image=image.convert('RGB')  #因为png是四个通道，除了RGB三个通道以外，还有一个透明度通道，加上这一步可以保留其颜色通道，可以适应jpg，png各种格式的图片
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
                                            torchvision.transforms.ToTensor()])  #转换成torch.Size([3, 32, 32])，输入到网络模型

image=transform(image)
print(image.shape)  #torch.Size([3, 32, 32])

class Aaax(nn.Module):  # 将 nn.module 改为 nn.Module
    def __init__(self):
        super(Aaax, self).__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, x):
        x = self.model(x)
        return x
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 检查GPU是否可用
model = torch.load("aa_9.pth", map_location=device)  # 加载模型并将其移动到指定设备   用gpu训练的模型就加载cuda的测试
model.to(device)  # 确保模型在设备上
print(model)  #模型已经加载好了

image=torch.reshape(image,(1,3,32,32)) #需要增加batch维度
image=image.to(device)


model.eval()
with torch.no_grad():
    output = model(image)
print(output)
print(output.argmax(1)[0])

使用之前训练的模型识别狗的图片：