diff --git a/test_IVF.py b/test_IVF.py index ef78c63..09ad867 100644 --- a/test_IVF.py +++ b/test_IVF.py @@ -17,14 +17,16 @@ current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") os.environ["CUDA_VISIBLE_DEVICES"] = "0" -ckpt_path= r"/home/star/whaiDir/PFCFuse/models/whaiFusion11-17-10-34.pth" +pth_path = "whaiFusion11-15-17-48" +ckpt_path= r"/home/star/whaiDir/PFCFuse/models/"+pth_path+".pth" +print("path_pth:{}".format(ckpt_path)) for dataset_name in ["sar"]: print("\n"*2+"="*80) - model_name="PFCFuse Enhance 增加widthblock" + model_name="PFCFuse 最基本版本 " print("The test result of "+dataset_name+' :') test_folder = os.path.join('test_img', dataset_name) - test_out_folder=os.path.join('test_result',current_time,dataset_name) + test_out_folder=os.path.join('test_result',pth_path,dataset_name) device = 'cuda' if torch.cuda.is_available() else 'cpu' Encoder = nn.DataParallel(Restormer_Encoder()).to(device) diff --git a/train.py b/train.py index 583fdb2..04537d5 100644 --- a/train.py +++ b/train.py @@ -262,5 +262,7 @@ if True: 'BaseFuseLayer': BaseFuseLayer.state_dict(), 'DetailFuseLayer': DetailFuseLayer.state_dict(), } - torch.save(checkpoint, os.path.join("models/whaiFusion"+timestamp+'.pth')) + savepth = os.path.join("models/whaiFusion" + timestamp + '.pth'); + torch.save(checkpoint, savepth) + print("save model:{}".format(savepth))