state_dict_hook might corrupt the model when saving

[arrufat]

model_preparation_algorithm/mpa/modules/models/classifiers/sam_classifier.py

Lines 106 to 137 in 55af1ed

	if backbone_type == 'OTEMobileNetV3':
	for k, v in state_dict.items():
	if k.startswith('backbone'):
	k = k.replace('backbone.', '')
	elif k.startswith('head'):
	k = k.replace('head.', '')
	if '3' in k: # MPA uses "classifier.3", OTE uses "classifier.4". Convert for OTE compatibility.
	k = k.replace('3', '4')
	if module.multilabel and not module.is_export:
	v = v.t()
	output[k] = v
	elif backbone_type == 'OTEEfficientNet':
	for k, v in state_dict.items():
	if k.startswith('backbone'):
	k = k.replace('backbone.', '')
	elif k.startswith('head'):
	k = k.replace('head', 'output')
	if not module.hierarchical and not module.is_export:
	k = k.replace('fc', 'asl')
	v = v.t()
	output[k] = v
	elif backbone_type == 'OTEEfficientNetV2':
	for k, v in state_dict.items():
	if k.startswith('backbone'):
	k = k.replace('backbone.', '')
	elif k == 'head.fc.weight':
	k = k.replace('head.fc', 'model.classifier')
	if not module.hierarchical and not module.is_export:
	v = v.t()
	output[k] = v

Here, we're looking for keys that start with backbone or head, and then erasing them.
The problem is that we are not only erasing the first one, but all instances.
I had a model with an auxiliary head, named aux_head so this happened:

head.aux_head.weight → aux_weight

We should use k.replace('backbone.', 1) and k.replace('head.', 1) to make sure we only replace the first appearance.

If at some point we update to Python 3.9 we could even use k = k.remove_prefix('backbone.').