whaifree/pfcfuse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/base' into base_增加WTconv

Browse Source 
# Conflicts:
#	test_IVF.py
This commit is contained in:
zjut 2024-11-18 09:39:41 +08:00
commit 8383c68ff7
2 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
test_IVF.py
View File

@ -17,14 +17,16 @@ current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
os.environ["CUDA_VISIBLE_DEVICES"] = "0" 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"]: for dataset_name in ["sar"]:
print("\n"*2+"="*80) print("\n"*2+"="*80)
model_name="PFCFuse Enhance 增加widthblock" model_name="PFCFuse 最基本版本 "
print("The test result of "+dataset_name+' :') print("The test result of "+dataset_name+' :')
test_folder = os.path.join('test_img', 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' device = 'cuda' if torch.cuda.is_available() else 'cpu'
Encoder = nn.DataParallel(Restormer_Encoder()).to(device) Encoder = nn.DataParallel(Restormer_Encoder()).to(device)

4
train.py
View File

@ -262,5 +262,7 @@ if True:
'BaseFuseLayer': BaseFuseLayer.state_dict(), 'BaseFuseLayer': BaseFuseLayer.state_dict(),
'DetailFuseLayer': DetailFuseLayer.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))