用pytorch构建Alexnet模型(predict模块)(个人笔记)

将原代码修改了一下,可以将所需要预测的图片放在指定文件夹内(可放多个),代码会依次预测该图片属于哪个类,并将其保存在TXT文件中(run栏不显示预测结果不要感到奇怪,打开当前目录下生成的txt文件即可看到结果)。

import os
import json

import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt

from model import AlexNet


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    data_transform = transforms.Compose(    # 定义图片预处理函数,用来对载入图片进行预处理操作
        [transforms.Resize((224, 224)),   # 缩放到224*224
         transforms.ToTensor(),   # 转化为一个tensor
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])    # 标准化处理
    txtfilepath = "F:\PyTorch\picture"   # 原始txt文件所存文件夹,文件夹可以有一个或多个txt文件
    total_txt = os.listdir(txtfilepath)   # 返回指定的文件夹包含的文件或文件夹的名字的列表
    num = len(total_txt)
    list = range(num)  # 创建从0到num的整数列表

    for i in list:
        name = total_txt[i]
        # load image
        img = Image.open(txtfilepath+"/"+name, 'r') #读取文件

        plt.imshow(img)  # 展示输入的图片
        # [N, C, H, W]
        img = data_transform(img)   # 调用预处理函数,对载入读片进行预处理
        # expand batch dimension
        img = torch.unsqueeze(img, dim=0)  # 预处理之后扩充一个维度(batch维度),这与Alexnet输入有关(具体见NB笔记)

        # read class_indict
        json_path = './class_indices.json'   # 读取保存的json文件(类别名称以及对应的索引)
        assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)

        json_file = open(json_path, "r")  # 解码成所需要的字典
        class_indict = json.load(json_file)

        # create model
        model = AlexNet(num_classes=5).to(device)   # 初始化网络

        # load model weights
        weights_path = "./AlexNet.pth"
        assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
        model.load_state_dict(torch.load(weights_path))  # 载入网络模型

        model.eval()   # 进入eval模式(即关闭掉droout方法)
        with torch.no_grad():   # 不跟踪变量的损失梯度
            # predict class
            output = torch.squeeze(model(img.to(device))).cpu()   # 将图片通过model正向传播,得到输出,将输入进行压缩,将batch维度压缩掉,得到最终输出(out)
            predict = torch.softmax(output, dim=0)  # 经过softmax处理后,就变成概率分布的形式了
            predict_cla = torch.argmax(predict).numpy()  # 通过argmax方法,得到概率最大的处所对应的索引


        print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
                                                 predict[predict_cla].numpy())   # 打印类别名称以及他所对应的预测概率

        plt.title(print_res)
        for i in range(len(predict)):
            with open('test.txt', 'a') as file0:  # 将以下print内容保存到test.txt文件中
                print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
                                                      predict[i].numpy()), file=file0)
        with open('test.txt', 'a') as file0:
            print("--------------------------我是可爱的分隔线--------------------------", file=file0)
        # plt.show()


if __name__ == '__main__':
    main()

你可能感兴趣的:(深度学习,算法,python)