pytorch-模型训练-加载图像分类模型训练参数

 将训练好的模型参数进行测试，检测其正确率，并且写入csv文件中，其每个类的概率值，本模型为图像分类模型：

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # 按照PCI_BUS_ID顺序从0开始排列GPU设备 
os.environ["CUDA_VISIBLE_DEVICES"] = "1" #设置当前使用的GPU设备仅为0号设备  设备名称为'/gpu:0'
import sys
from PIL import Image
import torchvision.transforms.functional as tv_F
from torch.autograd import Variable
import numpy as np
from models.config import cfg

def result_(res):
    if res==0:
                return 'airplane'
    elif res==1:
                return 'ship'
    elif res==2:
                return 'bridge'
    elif res==3:
                return 'oilcan'
    elif res==4:
                return 'build'
    else:
                return 'Nan'

if __name__=='__main__':
    # csvFile = open('stag_02_noise_all_1.csv', 'w', newline='')
    # csvFile = open("stag_02_qb.csv", "w")  # 创建csv文件,会有空行
    # writer = csv.writer(csvFile)  # 创建写的对象
    # 先写入columns_name
    # writer.writerow(["id", "label", "ship", "bridge", "airplane", "build", "oilcan"])  # 写入列的名称

    model = build_model(cfg, num_classes=5)
    model.fc = nn.Linear(512*4, 5)

    save_dir = "./models/new/"
 ######  model加载有两种方法，正常为第一个，直接加载，但是如果报错，选择第二种方法
#     model.load_state_dict(torch.load(os.path.join(save_dir, 'modeling_stag_new.dat')))
    model.load_state_dict({k.replace('module.',''):v for k,v in torch.load(os.path.join(save_dir, 'modeling_stag_new.dat')).items()})
    testdata_dir = 'stag_03_test'  # 0001_0.jpg
    model.eval()
    if torch.cuda.is_available() :
        model = torch.nn.DataParallel(model).cuda()
    img_list = os.listdir(testdata_dir)
    count = 0
    for img_name in img_list:
        print(img_name)
        img_ = (img_name.split('.')[0])
        id_ = (img_name.split('.')[0]).split('_')[1]
        print('Processing image: ' + img_)
        img = Image.open(os.path.join(testdata_dir, img_name))
        img = tv_F.to_tensor(tv_F.resize(img, (299, 299)))
        img = tv_F.normalize(img, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        img_input = Variable(torch.unsqueeze(img, 0))
        img_input = img_input.cuda()
        # print(img_input.shape)#torch.Size([1, 3, 299, 299])
        output = model(img_input)
#         print(output.shape)
        # 计算softmax，即该图片属于各类的概率
        probability = F.softmax(output, dim=1)
        value,index  = torch.max(probability.data, dim=1) # 找到最大概率对应的索引号，该图片即为该索引号对应的类别
#         print('probability', len(probability[0]))#([[0.0065, 0.7353, 0.0955, 0.0198, 0.1429]]
        class_index = result_(index)
        # print(class_index)
        probability = np.round(probability.cpu().detach().numpy(), 3)
        # writer.writerow([img_name, class_index, probability[0][1], probability[0][2], probability[0][0], probability[0][4],probability[0][3]])
        pred = (index[0]).cpu().numpy()
        if str(pred) == id_:
            count += 1
            # print('Right')
        print('Prediction number: %s'%(pred))
        print('target number: ' + id_)
    correct_rate = (count / len(img_list)) * 100
    print('准确率为：%d' % (correct_rate))
    # csvFile.close()

打印结果为： 

csv文件为：