facebookresearch / fvcore

Collection of common code that's shared among different research projects in FAIR computer vision team.
Apache License 2.0
2k stars 226 forks source link

Sort state_dict.keys() in _strip_prefix_if_present? #136

Open isdanni opened 1 year ago

isdanni commented 1 year ago

https://github.com/facebookresearch/fvcore/blob/0f2b23b6f93e36041d9a74764ee824541cf0a0e5/fvcore/common/checkpoint.py#L518

Was wondering why sorting the state_dict.keys() here when it is already a SortedDict? This will obviously change the order of the state_dict keys.

Example:

from collections import OrderedDict

state_dict = OrderedDict()
state_dict['module.body.stage1.0.0.weight'] = 0
state_dict['module.body.stage1.0.1.weight'] = 1
state_dict['module.body.stage1.0.1.bias'] = 2
state_dict['module.bboxes.0.conv1x1.weight'] = 3
state_dict['module.bboxes.0.conv1x1.bias'] = 4

keys_og = state_dict.keys()
"""
['module.body.stage1.0.0.weight', 'module.body.stage1.0.1.weight', 'module.body.stage1.0.1.bias', 'module.bboxes.0.conv1x1.weight', 'module.bboxes.0.conv1x1.bias']
"""
keys = sorted(state_dict.keys())
"""
['module.bboxes.0.conv1x1.bias', 'module.bboxes.0.conv1x1.weight', 'module.body.stage1.0.0.weight', 'module.body.stage1.0.1.bias', 'module.body.stage1.0.1.weight']
"""

prefix = "module."

for key in keys_og:
    if key.startswith(prefix):
        newkey = key[len(prefix) :]
        state_dict[newkey] = state_dict.pop(key)

"""
With sorting: 
OrderedDict([('bboxes.0.conv1x1.bias', 4), ('bboxes.0.conv1x1.weight', 3), ('body.stage1.0.0.weight', 0), ('body.stage1.0.1.bias', 2), ('body.stage1.0.1.weight', 1)])

Without sorting:
OrderedDict([('body.stage1.0.0.weight', 0), ('body.stage1.0.1.weight', 1), ('body.stage1.0.1.bias', 2), ('bboxes.0.conv1x1.weight', 3), ('bboxes.0.conv1x1.bias', 4)])
"""