(
self,
model_encoder: nn.Module,
model_task: nn.Module,
optim_cfg: partial,
scheduler_cfg: Optional[partial],
checkpoint_encoder: Optional[str] = None,
checkpoint_task: Optional[str] = None,
train_encoder: bool = False,
encoder_type: str = "jepa",
is_anytouch: bool = False,
)
| 17 | |
| 18 | class SLModule(Module, nn.Module): |
| 19 | def __init__( |
| 20 | self, |
| 21 | model_encoder: nn.Module, |
| 22 | model_task: nn.Module, |
| 23 | optim_cfg: partial, |
| 24 | scheduler_cfg: Optional[partial], |
| 25 | checkpoint_encoder: Optional[str] = None, |
| 26 | checkpoint_task: Optional[str] = None, |
| 27 | train_encoder: bool = False, |
| 28 | encoder_type: str = "jepa", |
| 29 | is_anytouch: bool = False, |
| 30 | ) -> None: |
| 31 | super().__init__() |
| 32 | self.model_task: nn.Module = model_task |
| 33 | self.model_encoder: nn.Module = model_encoder |
| 34 | self.train_encoder: bool = train_encoder |
| 35 | self.encoder_type: str = encoder_type |
| 36 | |
| 37 | if checkpoint_encoder is not None: |
| 38 | log.info("Loading encoder ONLY from checkpoint.") |
| 39 | self.load_encoder(checkpoint_encoder) |
| 40 | else: |
| 41 | log.info("No checkpoint provided. Training from scratch.") |
| 42 | |
| 43 | if checkpoint_task is not None: |
| 44 | log.info("Loading task decoder from checkpoint.") |
| 45 | self.load_task(checkpoint_task) |
| 46 | |
| 47 | # freeze encoder |
| 48 | if not self.train_encoder: |
| 49 | self.model_encoder.requires_grad_(False) |
| 50 | self.model_encoder.eval() |
| 51 | self.scheduler_partial = scheduler_cfg |
| 52 | self.optim_partial = optim_cfg |
| 53 | |
| 54 | def load_task(self, checkpoint_task: str): |
| 55 | try: |
nothing calls this directly
no test coverage detected