Coverage for /scratch/builds/bob/bob.med.tb/miniconda/conda-bld/bob.med.tb_1637571489937/_test_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placeho/lib/python3.8/site-packages/bob/med/tb/utils/model_serialization.py: 0%

1# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 

2# https://github.com/facebookresearch/maskrcnn-benchmark 

3 

4from collections import OrderedDict 

5 

6import logging 

7logger = logging.getLogger(__name__) 

8 

9import torch 

10 

11 

12def align_and_update_state_dicts(model_state_dict, loaded_state_dict): 

13 """ 

14 Strategy: suppose that the models that we will create will have prefixes appended 

15 to each of its keys, for example due to an extra level of nesting that the original 

16 pre-trained weights from ImageNet won't contain. For example, model.state_dict() 

17 might return backbone[0].body.res2.conv1.weight, while the pre-trained model contains 

18 res2.conv1.weight. We thus want to match both parameters together. 

19 For that, we look for each model weight, look among all loaded keys if there is one 

20 that is a suffix of the current weight name, and use it if that's the case. 

21 If multiple matches exist, take the one with longest size 

22 of the corresponding name. For example, for the same model as before, the pretrained 

23 weight file can contain both res2.conv1.weight, as well as conv1.weight. In this case, 

24 we want to match backbone[0].body.conv1.weight to conv1.weight, and 

25 backbone[0].body.res2.conv1.weight to res2.conv1.weight. 

26 """ 

27 current_keys = sorted(list(model_state_dict.keys())) 

28 loaded_keys = sorted(list(loaded_state_dict.keys())) 

29 # get a matrix of string matches, where each (i, j) entry correspond to the size of the 

30 # loaded_key string, if it matches 

31 match_matrix = [ 

32 len(j) if i.endswith(j) else 0 for i in current_keys for j in loaded_keys 

33 ] 

34 match_matrix = torch.as_tensor(match_matrix).view( 

35 len(current_keys), len(loaded_keys) 

36 ) 

37 max_match_size, idxs = match_matrix.max(1) 

38 # remove indices that correspond to no-match 

39 idxs[max_match_size == 0] = -1 

40 

41 # used for logging 

42 max_size = max([len(key) for key in current_keys]) if current_keys else 1 

43 max_size_loaded = max([len(key) for key in loaded_keys]) if loaded_keys else 1 

44 log_str_template = "{: <{}} loaded from {: <{}} of shape {}" 

45 for idx_new, idx_old in enumerate(idxs.tolist()): 

46 if idx_old == -1: 

47 continue 

48 key = current_keys[idx_new] 

49 key_old = loaded_keys[idx_old] 

50 model_state_dict[key] = loaded_state_dict[key_old] 

51 logger.debug( 

52 log_str_template.format( 

53 key, 

54 max_size, 

55 key_old, 

56 max_size_loaded, 

57 tuple(loaded_state_dict[key_old].shape), 

58 ) 

59 ) 

60 

61 

62def strip_prefix_if_present(state_dict, prefix): 

63 keys = sorted(state_dict.keys()) 

64 if not all(key.startswith(prefix) for key in keys): 

65 return state_dict 

66 stripped_state_dict = OrderedDict() 

67 for key, value in state_dict.items(): 

68 stripped_state_dict[key.replace(prefix, "")] = value 

69 return stripped_state_dict 

70 

71 

72def load_state_dict(model, loaded_state_dict): 

73 model_state_dict = model.state_dict() 

74 # if the state_dict comes from a model that was wrapped in a 

75 # DataParallel or DistributedDataParallel during serialization, 

76 # remove the "module" prefix before performing the matching 

77 loaded_state_dict = strip_prefix_if_present(loaded_state_dict, prefix="module.") 

78 align_and_update_state_dicts(model_state_dict, loaded_state_dict) 

79 

80 # use strict loading 

81 model.load_state_dict(model_state_dict)