Pytorch模型的推理如何编写-以猫狗分类为例

📚博客主页：knighthood2001
✨公众号：认知up吧 （目前正在带领大家一起提升认知，感兴趣可以来围观一下）
🎃知识星球：【认知up吧|成长|副业】介绍
❤️如遇文章付费，可先看看我公众号中是否发布免费文章❤️
🙏笔者水平有限，欢迎各位大佬指点，相互学习进步！

全部代码

import torch
import torch.utils.data as Data
from torchvision import transforms
from model import GoogLeNet, Inception
from torchvision.datasets import ImageFolder
from PIL import Image

def test_model(model, test_file):
    # 设定测试所用到的设备，有GPU用GPU没有GPU用CPU
    device = "cuda" if torch.cuda.is_available() else 'cpu'
    model = model.to(device)
    classes = ['猫', '狗']
    print(classes)
    image = Image.open(test_file)

    # normalize = transforms.Normalize([0.162, 0.151, 0.138], [0.058, 0.052, 0.048])
    # # 定义数据集处理方法变量
    # test_transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor(), normalize])

    # 定义数据集处理方法变量
    test_transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
    image = test_transform(image)

    # 添加批次维度，变成[1,3,224,224]
    image = image.unsqueeze(0)

    with torch.no_grad():
        model.eval()
        image = image.to(device)  # 图片也要放到设备当中
        output = model(image)
        print(output.tolist())
        pre_lab = torch.argmax(output, dim=1)
        result = pre_lab.item()
    print("预测值：", classes[result])

if __name__ == "__main__":
    # 加载模型
    model = GoogLeNet(Inception, in_channels=3, out_channels=2)
    model.load_state_dict(torch.load('best_model.pth'))

    # 模型的推理判断
    test_model(model, "test_data/images.jfif")

讲解

首先，你需要模型，训练好的模型参数，以及你的测试文件。

• 定义好你的设别，cpu还是gpu
• 模型放到设备上
• 打开你的测试文件，进行数据预处理
• 对数据补维度（看情况）
• 将模型设置为评估模式
• 将测试文件也放到设备上
• 通过模型产生结果

运行结果