Problem about testing my trained model
hcleung3325 opened this issue · comments
Thanks for the training code.
I train a classic model.
I get 500000_optimizerG.pth 500000_G.pth 500000_E.pth.
May I know which pth I should run during testing?
When I run
python main_test_swinir.py --task classical_sr --scale 4 --training_patch_size 64 --model_path superresolution/swinir_sr_classical_patch64_x4_l1/models/500000_G.pth --folder_lq testsets/real3wx4/test_LR_crop
It seems cannot load the model
model.load_state_dict(torch.load(args.model_path)['params'], strict=True)
KeyError: 'params'
May I know how to solve this?
Thanks.
Hi, sorry for this bug. Please remove ['params']
in Line 126
Line 126 in 96a446f
as
model.load_state_dict(torch.load(args.model_path), strict=True).
Hi, sorry for this bug. Please remove
['params']
in Line 126
Line 126 in 96a446f
as
model.load_state_dict(torch.load(args.model_path), strict=True).
Thanks for reply.
I notice that when I ran your pretrained model, the main_test_swinir.py works fine.
However, when I ran main_test_swinir.py with either 500000_optimizerG.pth 500000_G.pth 500000_E.pth, it occurs the above error.
Is it the 500000_G.pth has some problem?
The main difference between my pretrained model and yours (trained on KAIR) is that:
1, my pretrained model is a dict: {'params': model}
.
2, KAIR's model is just the model
.
Therefore, you have to use the following line to load my pretrained model:
model.load_state_dict(torch.load(args.model_path)['params'], strict=True) # specify the dict key
and the following line to load your trained model (on KAIR)
model.load_state_dict(torch.load(args.model_path), strict=True) # directly load it
Sorry for the inconsistence between KAIR and SwinIR testing code. We will fix this problem.
Update: I have fixed the bug.
Thanks a lot.