(self, checkpoint_encoder: str)
| 80 | log.info(f"Could not load task model from {checkpoint_task}") |
| 81 | |
| 82 | def load_encoder(self, checkpoint_encoder: str): |
| 83 | log.info(f"Loading encoder from {checkpoint_encoder}") |
| 84 | checkpoint = torch.load(checkpoint_encoder) |
| 85 | if "jepa" in self.encoder_type: |
| 86 | encoder_key = "target_encoder" |
| 87 | elif "dino" in self.encoder_type: |
| 88 | encoder_key = "teacher_encoder.backbone" |
| 89 | else: |
| 90 | encoder_key = "encoder" |
| 91 | # get the keys in the checkpoint that contain the encoder |
| 92 | target_keys = [key for key in checkpoint["model"].keys() if encoder_key in key] |
| 93 | if 'backbone' in target_keys[0] and 'backbone' not in encoder_key: |
| 94 | encoder_key = encoder_key + '.backbone' |
| 95 | # remove the prefix from the keys |
| 96 | new_keys = [key.replace(f"{encoder_key}.", "") for key in target_keys] |
| 97 | # create a state_dict with keys target_keys from the checkpoint |
| 98 | new_state_dict = { |
| 99 | new_key: checkpoint["model"][target_key] |
| 100 | for new_key, target_key in zip(new_keys, target_keys) |
| 101 | } |
| 102 | # load the state_dict into the model |
| 103 | self.model_encoder.load_state_dict(new_state_dict, strict=False) |
| 104 | log.info(f"Loaded encoder from {checkpoint_encoder}") |
| 105 | |
| 106 | def forward(self, x, *args, **kwargs): # noqa |
| 107 | raise NotImplementedError |
no test coverage detected