mmpretrain [Bug] Revise the _remove_state_dict_prefix and _add_state_dict_prefix functions in timm.py to adapt to the case of multiple submodels.

[Bug] Revise the _remove_state_dict_prefix and _add_state_dict_prefix functions in timm.py to adapt to the case of multiple submodels.

Open wilxy opened this issue 2 years ago • 5 comments

When using TimmClassifier as student or teacher model in Knowledge Distillation Algorithms, there have some bugs in save_checkpoint and load_checkpoint.

save_checkpoint When saving checkpoint like save_checkpoint(self.state_dict(), 'xxx.pth'), where self is a Knowledge Distillation Algorithm which contains submodels self.student and self.teacher, self.state_dict() will recursively call the state_dict function here. The _remove_state_dict_prefix function in the TimmClassifier class will be used as a hook to modify the original destination. Specifically, the _remove_state_dict_prefix function creates a new_state_dict whose memory is different from the original destination as the hook_result to modify the original destination for submodels student and teacher. But the state_dict funtion of the Knowledge Distillation Algorithm Model will not receive this modify, so the memory address and value of destination have not changed. To solve this problem, we change the _remove_state_dict_prefix function to modify the state_dict directly instead of creating a new_state_dict.
load_checkpoint When loading checkpoint of a Knowledge Distillation Algorithm Model whose student and teacher are all TimmClassifier. The _add_state_dict_prefix function in the TimmClassifier class will be used as a hook to modify the state_dict of each submodel. When modifying the student submodel, _add_state_dict_prefix function will delete all keys of teacher submodel. To solve this problem, we change the _add_state_dict_prefix function to only delete the key that different from its new_key.