(config, model, logger)
| 101 | |
| 102 | |
| 103 | def load_pretrained(config, model, logger): |
| 104 | logger.info( |
| 105 | f'==============> Loading weight {config.MODEL.PRETRAINED} for fine-tuning......' |
| 106 | ) |
| 107 | checkpoint = torch.load(config.MODEL.PRETRAINED, map_location='cpu') |
| 108 | |
| 109 | state_dict = checkpoint |
| 110 | if 'model' in checkpoint: |
| 111 | state_dict = checkpoint['model'] |
| 112 | elif 'module' in checkpoint: |
| 113 | state_dict = checkpoint['module'] |
| 114 | |
| 115 | first_key = list(state_dict.keys())[0] |
| 116 | # delete teacher weights |
| 117 | if 'student' in first_key or 'teacher' in first_key: |
| 118 | new_state_dict = OrderedDict() |
| 119 | for k, v in state_dict.items(): |
| 120 | if 'student_proj' in k: |
| 121 | continue |
| 122 | if 'student' in k: |
| 123 | new_k = k.replace('student.', '') |
| 124 | new_state_dict[new_k] = v |
| 125 | state_dict = new_state_dict |
| 126 | |
| 127 | # weights from sim |
| 128 | if 'mask_token' in first_key: |
| 129 | new_state_dict = OrderedDict() |
| 130 | for k, v in state_dict.items(): |
| 131 | if 'mm_dcnv3' in k: |
| 132 | continue |
| 133 | if 'dcnv3' not in k and 'clip_projector' not in k: |
| 134 | continue |
| 135 | new_k = k.replace('dcnv3.', '') |
| 136 | new_state_dict[new_k] = v |
| 137 | new_state_dict['fc_norm.weight'] = state_dict[ |
| 138 | 'clip.classifier_ln.weight'] |
| 139 | new_state_dict['fc_norm.bias'] = state_dict['clip.classifier_ln.bias'] |
| 140 | new_state_dict['head.weight'] = state_dict['clip.classifier.weight'] |
| 141 | new_state_dict['head.bias'] = state_dict['clip.classifier.bias'] |
| 142 | state_dict = new_state_dict |
| 143 | |
| 144 | # delete relative_position_index since we always re-init it |
| 145 | relative_position_index_keys = [ |
| 146 | k for k in state_dict.keys() if 'relative_position_index' in k |
| 147 | ] |
| 148 | for k in relative_position_index_keys: |
| 149 | del state_dict[k] |
| 150 | |
| 151 | # delete relative_coords_table since we always re-init it |
| 152 | relative_position_index_keys = [ |
| 153 | k for k in state_dict.keys() if 'relative_coords_table' in k |
| 154 | ] |
| 155 | for k in relative_position_index_keys: |
| 156 | del state_dict[k] |
| 157 | |
| 158 | # delete attn_mask since we always re-init it |
| 159 | attn_mask_keys = [k for k in state_dict.keys() if 'attn_mask' in k] |
| 160 | for k in attn_mask_keys: |
no test coverage detected