(self, checkpoint_task: str)
| 52 | self.optim_partial = optim_cfg |
| 53 | |
| 54 | def load_task(self, checkpoint_task: str): |
| 55 | try: |
| 56 | state_dict = torch.load(checkpoint_task) |
| 57 | # check if there are keys starting with "model_encoder." |
| 58 | if any([key.startswith("model_encoder.") for key in state_dict.keys()]): |
| 59 | log.info("Found encoder in task checkpoint. Loading encoder and decoder from task checkpoint.") |
| 60 | self.load_state_dict(state_dict, strict=False) |
| 61 | return |
| 62 | else: |
| 63 | state_dict = { |
| 64 | key.replace("model_task.", ""): value |
| 65 | for key, value in state_dict.items() |
| 66 | } |
| 67 | self.model_task.load_state_dict(state_dict, strict=False) |
| 68 | log.info(f"Loaded task model from {checkpoint_task}") |
| 69 | except: |
| 70 | # add to state_dict_light only keys that start with model_encoder |
| 71 | try: |
| 72 | state_dict_light = { |
| 73 | key.replace("model_encoder.", ""): value |
| 74 | for key, value in torch.load(checkpoint_task).items() |
| 75 | if key.startswith("model_encoder.") |
| 76 | } |
| 77 | self.model_encoder.load_state_dict(state_dict_light, strict=False) |
| 78 | log.info(f"Loaded encoder from {checkpoint_task}") |
| 79 | except: |
| 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}") |
no test coverage detected