JingyunLiang / SwinIR

SwinIR: Image Restoration Using Swin Transformer (official repository)

Home Page:https://arxiv.org/abs/2108.10257

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

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

model.load_state_dict(torch.load(args.model_path)['params'], strict=True)

as

model.load_state_dict(torch.load(args.model_path), strict=True). 

Hi, sorry for this bug. Please remove ['params'] in Line 126

model.load_state_dict(torch.load(args.model_path)['params'], strict=True)

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.