（二）pytorch利用训练好的模型来进行预测

前面我们训练的用的数据是ants(蚂蚁)和bees(蜜蜂), 所以预测时用的图片应该跟训练时类似。
本次测试的图片为ants，如下图所示：

代码中的图片路径和模型路径根据自己的情况来修改
预测代码为：

import torch, glob, cv2
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
import matplotlib.pyplot as plt


def preict_one_img(img_path, model_path):
    net = torch.load(model_path)
    img = cv2.imread(img_path)
    img = cv2.resize(img, (224, 224))
    # 1.将numpy数据变成tensor
    tran = transforms.ToTensor()
    img = tran(img)
    img = img.to(device)
    # 2.将数据变成网络需要的shape
    img = img.view(1, 3, 224, 224)

    out1 = net(img)
    out1 = F.softmax(out1, dim=1)
    proba, class_ind = torch.max(out1, 1)

    proba = float(proba)
    class_ind = int(class_ind)
    # print(proba, class_ind)
    img = img.cpu().numpy().squeeze(0)
    # print(img.shape)
    new_img = np.transpose(img, (1,2,0))
    plt.imshow(new_img)
    plt.title("the predict is %s . prob is %s" % (classes[class_ind], round(proba, 3)))   # round(proba, 3)保留三位小数
    plt.show()



if __name__ == '__main__':
    classes = ["ants", "bees"]
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    img_path = r"E:\myworld\my_like\using_torch\project1_classify\hymenoptera_data\test\154124431_65460430f2.jpg"
    model_path = r"E:\myworld\my_like\using_torch\project1_classify\my_mdel.pth"
    preict_one_img(img_path, model_path)