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檔案為：