MogaNet实战指南：利用MogaNet进行图片识别任务（第二部分）- 测试篇

最编程 2024-07-26 19:12:29

...

测试，我们采用一种通用的方式。

测试集存放的目录如下图：

MogaNet_Demo
├─test
│  ├─1.jpg
│  ├─2.jpg
│  ├─3.jpg
│  ├ ......
└─test.py

import torch.utils.data.distributed
import torchvision.transforms as transforms
from PIL import Image
from torch.autograd import Variable
import os

classes = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed',
           'Common wheat', 'Fat Hen', 'Loose Silky-bent',
           'Maize', 'Scentless Mayweed', 'Shepherds Purse', 'Small-flowered Cranesbill', 'Sugar beet')
transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.44127703, 0.4712498, 0.43714803], std=[0.18507297, 0.18050247, 0.16784933])
])

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model=torch.load('checkpoints/MogaNet/best.pth')
model.eval()
model.to(DEVICE)
print(model)

path = 'test/'
testList = os.listdir(path)
for file in testList:
    img = Image.open(path + file)
    img = transform_test(img)
    img.unsqueeze_(0)
    img = Variable(img).to(DEVICE)
    out = model(img)
    # Predict
    _, pred = torch.max(out.data, 1)
    print('Image Name:{},predict:{}'.format(file, classes[pred.data.item()]))

测试的主要逻辑：

1、定义类别，这个类别的顺序和训练时的类别顺序对应，一定不要改变顺序！！！！

2、定义transforms，transforms和验证集的transforms一样即可，别做数据增强。

3、 torch.load加载model，然后将模型放在DEVICE里，

4、循环读取图片并预测图片的类别，在这里注意，读取图片用PIL库的Image。不要用cv2，transforms不支持。循环里面的主要逻辑：

使用Image.open读取图片

使用transform_test对图片做归一化和标椎化。

img.unsqueeze_(0) 增加一个维度，由（3，224，224）变为（1，3，224，224）

Variable(img).to(DEVICE)：将数据放入DEVICE中。

model(img)：执行预测。

_, pred = torch.max(out.data, 1):获取预测值的最大下角标。

运行结果：

在这里插入图片描述

上一篇：在PMP®考试培训里，工作分解结构(WBS)的实战意义是什么？

下一篇：轻松上手NCCL(初级篇)：简介与简易演示 - NCCL的安装指南