深度模型测试

import torch,gc
import sys
from torch.utils.data import DataLoader
from datasets_ import Drishti_IMG,Test_Drishti
from models_.net import UNet
from torchvision.utils import save_image
import os
from models_.myNet import MyNet
from Metrics import IOU_acc,Dice
from utils_.Util_methods import  AddInterpreter


""" 四个网络 """


if __name__ == '__main__':

    # 设置数据集
    data_root=os.path.join("/home/wzc/zlt_self/pOSAL-master","data")
    if os.path.exists(data_root):
        dataloader=DataLoader(dataset=Test_Drishti(data_root),batch_size=4,shuffle=False)
    else:
        print("{} 文件夹不存在".format(data_root))
        
    # model加载到GPU，然后加载权重参数，最后设置优化和损失函数
    device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model_=UNet().to(device=device)
    for index_path in range(4):
        weight_path=f'/home/wzc/zlt_self/pOSAL-master/LarryChow/draft/tempFiles/weights_/unet_{index_path}.pth'
        # weight_path='/home/wzc/zlt_self/pOSAL-master/LarryChow/draft/tempFiles/weights_/unet.pth'
        if os.path.exists(weight_path):
            model_.load_state_dict(torch.load(weight_path))
            print('successful load weight！')
            
        else:
            print("not successful load weight")
            sys.exit(0)
        
        # 设置损失函数和精确度
        dice_fun=Dice()
        acc_fun=IOU_acc()
        
        for i,(img_,label_) in enumerate(dataloader):
            if i==1:
                continue
            #释放内存
            gc.collect()
            torch.cuda.empty_cache()
            # 模型预测
            img_=img_.to(device)
            label_=label_.to(device)
            pred_=model_(img_)  #type(pred_)=<class 'torch.Tensor'>
            for x_ in range(4):
                for index_ in range(1):
                    # save_img_=torch.stack([img_[x_],label_[x_],pred_[x_]],dim=0)
                    save_image(pred_[x_],f"unet_{index_path}-{x_}.png")
                    #添加文字信息
                    # AddInterpreter(imgPath=f"/home/wzc/zlt_self/pOSAL-master/tempFiles/unet-test_result{i}-{x_}.png",savePath=f"/home/wzc/zlt_self/pOSAL-master/tempFiles/unet-test_result_{i}-{x_}.png")
                    # print("pred success...")
                    # # 损失函数和精确度
                    # test_dice=dice_fun(label_,pred_)
                    # test_acc=acc_fun(label_,pred_)
                    # print(f"***{i}-{x_}test_dice:{test_dice}   test_acc:{test_acc}\n")
                    print(f"{index_path}{i}-{x_}-{index_}OK")
            break