Merge pull request #5 from carstenli/patch-3

build_backbone.py update to make it work on Apple Silicon devices
This commit is contained in:
Xu Yang
2024-03-25 18:49:36 +08:00
committed by GitHub
+2 -2
View File
@@ -26,7 +26,7 @@ def build_backbone(bb_name, pretrained=True, params_settings=''):
return bb
def load_weights(model, model_name):
save_model = torch.load(config.weights[model_name])
save_model = torch.load(config.weights[model_name], map_location=torch.device('cpu'))
model_dict = model.state_dict()
state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()}
# to ignore the weights with mismatched size when I modify the backbone itself.
@@ -35,7 +35,7 @@ def load_weights(model, model_name):
sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None
state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()}
if not state_dict or not sub_item:
print('Weights are not successully loaded. Check the state dict of weights file.')
print('Weights are not successfully loaded. Check the state dict of weights file.')
return None
else:
print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item))