pytorch-将网络模型预测结果的概率值存放到csv文件中

import csv
import cv2
#stag_01_submit
csvFile = open("test_submit.csv", "w")            #创建csv文件
writer = csv.writer(csvFile)                  #创建写的对象
#先写入columns_name     
writer.writerow(["id","label","ship","bridge","airplane","build","oilcan"])     #写入列的名称

#读入图片
test_root = './stag_04/'
img_test=os.listdir(test_root)

for i in range(len(img_test)): 
    rd_img = cv2.imread(test_root+img_test[i])
    img = Image.open(test_root+img_test[i]) 
    img = img.convert('RGB')
  
    input=trans(img)      
    input=input.unsqueeze(0)#这里经过转换后输出的input格式是[C,H,W],网络输入还需要增加一维批量大小B
                                      #增加一维，输出的img格式为[1,C,H,W]
    input = V(input.cuda())
    score = model(input)#将图片输入网络得到输出
    probability = t.nn.functional.softmax(score,dim=1)#计算softmax，即该图片属于各类的概率
    max_value,index = t.max(probability,1)#找到最大概率对应的索引号，该图片即为该索引号对应的类别
    class_index = result_(index)
    probability=np.round(probability.cpu().detach().numpy(),3)
    writer.writerow([img_test[i],class_index,probability[0][1],probability[0][2],probability[0][0],probability[0][4],probability[0][3]])
csvFile.close()

结果为：