From 64942273acab8e1b1cafa9af4dcf38f268ddddc5 Mon Sep 17 00:00:00 2001 From: Andrey Gurevich Date: Thu, 14 Jul 2022 12:01:39 +0300 Subject: [PATCH] Update helpers.py Should Fix errors like RuntimeError: Error(s) in loading state_dict for EfficientDet: Missing key(s) in state_dict: "backbone.conv_stem.weight", "backbone.bn1.weight", ... Unexpected key(s) in state_dict: "model.model.backbone.conv_stem.weight", "model.model.backbone.bn1.weight", "model.model.backbone.bn1.bias", "model.model.backbone.bn1.running_mean", "model.model.backbone.bn1.running_var", "model.model.backbone.bn1.num_batches_tracked", "model.model.backbone.blocks.0.0.conv_dw.weight", --- timm/models/helpers.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/timm/models/helpers.py b/timm/models/helpers.py index fda84171..9734c779 100644 --- a/timm/models/helpers.py +++ b/timm/models/helpers.py @@ -37,8 +37,14 @@ def clean_state_dict(state_dict): # 'clean' checkpoint by removing .module prefix from state dict if it exists from parallel training cleaned_state_dict = OrderedDict() for k, v in state_dict.items(): - name = k[7:] if k.startswith('module.') else k - cleaned_state_dict[name] = v + # strip `module.` and `model.` prefixes + if k.startswith('module'): + name = k[7:] + elif k.startswith('model'): + name = k[6:] + else: + name = k + new_state_dict[name] = v return cleaned_state_dict