小土堆pytorch教程学习笔记P32

Happy_Cabbage

已于 2022-04-18 09:18:54 修改

阅读量1.2k

点赞数

分类专栏：小土堆pytorch教程学习笔记文章标签： pytorch

于 2022-04-10 20:51:56 首次发布

本文链接：https://blog.csdn.net/kchenlyee/article/details/124082298

版权

小土堆pytorch教程学习笔记专栏收录该内容

22 篇文章 19 订阅

订阅专栏

P32.完整的模型验证套路

测试、demo

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

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

image_path = "imgs/truck.jpeg"
image = Image.open(image_path)
print(image)
image = image.convert('RGB')

transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
                                            torchvision.transforms.ToTensor()])

image = transform(image)

class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, 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),
            Flatten(),
            nn.Linear(64*4*4, 64),
            nn.Linear(64, 10)
        )
    def forward(self, x):
        x = self.model(x)
        return x

model = torch.load("tudui_29_gpu.pth")
print(model)
image = torch.reshape(image, (1, 3, 32, 32))
model.eval()
with torch.no_grad():
    output = model(image)
print(output)

print(output.argmax(1))