search

zsyOAOA / DifFace

DifFace: Blind Face Restoration with Diffused Error Contraction (TPAMI, 2024)
Other
628 stars 42 forks source link

Model weight preservation in the Train diffused estimator (SwinIR) stage. #17

Open 24wenjie-li opened 11 months ago

24wenjie-li commented 11 months ago

I retrained the Train diffused estimator (SwinIR) stage and found that the saved weights are dict_keys(['iter_start','log_step','log_step_img','state_dict']) and not SwinIR's weights, so the inference reports a load model error. Would like to ask if the code in main_sr.py about saving model weights is wrong? I'm looking forward to hearing from the author!

zsyOAOA commented 11 months ago

You should load the key of "state_dict". I have extracted this key for the released model weight.

24wenjie-li commented 11 months ago

You should load the key of "state_dict". I have extracted this key for the released model weight.

Thanks for the tip!