| 144 | params.requires_grad = found_name |
| 145 | |
| 146 | def load_state_with_same_shape(model, weights, prefix=''): |
| 147 | |
| 148 | model_state = model.state_dict() |
| 149 | if list(weights.keys())[0].startswith('module.'): |
| 150 | logging.info("Loading multigpu weights with module. prefix...") |
| 151 | weights = {k.partition('module.')[2]: weights[k] for k in weights.keys()} |
| 152 | |
| 153 | if list(weights.keys())[0].startswith('model.'): |
| 154 | logging.info("Loading Pytorch-Lightning weights from state") |
| 155 | weights = {k.partition('model.')[2]: weights[k] for k in weights.keys()} |
| 156 | |
| 157 | if list(weights.keys())[0].startswith('encoder.'): |
| 158 | logging.info("Loading multigpu weights with encoder. prefix...") |
| 159 | weights = {k.partition('encoder.')[2]: weights[k] for k in weights.keys()} |
| 160 | |
| 161 | if prefix != '': |
| 162 | weights = {k.partition(prefix)[2]: weights[k] for k in weights.keys()} |
| 163 | |
| 164 | # This is when checkpoint containes the full plmodule |
| 165 | target_keys = list(model_state.keys()) |
| 166 | if any(key.startswith('model_3d.') for key in target_keys) and any(key.startswith('model_2d.') for key in target_keys): |
| 167 | logging.info("Loading full lightning module weights") |
| 168 | bacbkbone_weights = np.array([key in target_keys for key in list(weights.keys())]) |
| 169 | bacbkbone_weights = np.array(list(weights.keys()))[bacbkbone_weights] |
| 170 | weights = {k: weights[k] for k in bacbkbone_weights} |
| 171 | |
| 172 | # For continuous if keys containing model_3d - this is when loading 3d model weights only from full checkpoint |
| 173 | if any(key.startswith('model_3d.') for key in list(weights.keys())) and not any(key.startswith('model_2d.') for key in target_keys): |
| 174 | logging.info("Loading backbone weights starting with model_3d.") |
| 175 | bacbkbone_weights = np.array([key.startswith('model_3d.') for key in list(weights.keys())]) |
| 176 | bacbkbone_weights = np.array(list(weights.keys()))[bacbkbone_weights] |
| 177 | weights = {k.partition('model_3d.')[2]: weights[k] for k in bacbkbone_weights} |
| 178 | |
| 179 | # print(weights.items()) |
| 180 | # print("===================") |
| 181 | # print("===================") |
| 182 | # print("===================") |
| 183 | # print("===================") |
| 184 | # print("===================") |
| 185 | # print(model_state) |
| 186 | |
| 187 | filtered_weights = { |
| 188 | k: v for k, v in weights.items() if k in model_state and v.size() == model_state[k].size() |
| 189 | } |
| 190 | logging.info(f"Loading weights for {len(filtered_weights.keys())}/{len(model_state.keys())} layers") |
| 191 | |
| 192 | return filtered_weights |
| 193 | |
| 194 | def load_matrix_from_txt(path, shape=(4, 4)): |
| 195 | with open(path) as f: |