(
self,
encoder: nn.Module,
dino_head: partial,
optim_cfg: partial,
lr_scheduler_cfg: Optional[partial],
wd_scheduler_cfg: Optional[partial],
online_probes: Optional[List[nn.Module]] = None,
online_probes_lrs: List[float] = [],
local_mask_scale: Tuple[float, float] = (0.2, 0.8),
global_mask_scale: Tuple[float, float] = (0.2, 0.8),
num_global_masks: int = 1,
num_local_masks: int = 4,
min_keep_num_sensors: int = 4,
allow_mask_overlap: bool = False,
moving_average_decay: Union[float, Tuple[float, ...]] = 0.99,
teacher_temp: Union[float, Tuple[float, ...]] = (0.04, 0.07),
teacher_warmup_epochs: int = 10,
use_momentum=True,
centering: Literal['centering', 'sinkhorn_knopp'] = 'centering',
ibot_separate_head: bool = False,
koleo_weight: float= 0.1,
log_freq_reconstruction: int = 1000,
)
| 30 | |
| 31 | class DINOv2Module(Module, nn.Module): |
| 32 | def __init__( |
| 33 | self, |
| 34 | encoder: nn.Module, |
| 35 | dino_head: partial, |
| 36 | optim_cfg: partial, |
| 37 | lr_scheduler_cfg: Optional[partial], |
| 38 | wd_scheduler_cfg: Optional[partial], |
| 39 | online_probes: Optional[List[nn.Module]] = None, |
| 40 | online_probes_lrs: List[float] = [], |
| 41 | local_mask_scale: Tuple[float, float] = (0.2, 0.8), |
| 42 | global_mask_scale: Tuple[float, float] = (0.2, 0.8), |
| 43 | num_global_masks: int = 1, |
| 44 | num_local_masks: int = 4, |
| 45 | min_keep_num_sensors: int = 4, |
| 46 | allow_mask_overlap: bool = False, |
| 47 | moving_average_decay: Union[float, Tuple[float, ...]] = 0.99, |
| 48 | teacher_temp: Union[float, Tuple[float, ...]] = (0.04, 0.07), |
| 49 | teacher_warmup_epochs: int = 10, |
| 50 | use_momentum=True, |
| 51 | centering: Literal['centering', 'sinkhorn_knopp'] = 'centering', |
| 52 | ibot_separate_head: bool = False, |
| 53 | koleo_weight: float= 0.1, |
| 54 | log_freq_reconstruction: int = 1000, |
| 55 | ): |
| 56 | super().__init__() |
| 57 | self.optim_partial = optim_cfg |
| 58 | self.lr_scheduler_partial = lr_scheduler_cfg |
| 59 | self.wd_scheduler_partial = wd_scheduler_cfg |
| 60 | self.use_momentum = use_momentum |
| 61 | self.global_mask_scale = global_mask_scale |
| 62 | self.local_mask_scale = local_mask_scale |
| 63 | self.num_global_masks = num_global_masks |
| 64 | self.num_local_masks = num_local_masks |
| 65 | self.min_keep = min_keep_num_sensors |
| 66 | self.centering = centering |
| 67 | self.allow_mask_overlap = allow_mask_overlap |
| 68 | self.log_freq_img = log_freq_reconstruction |
| 69 | self.ibot_separate_head = ibot_separate_head |
| 70 | self.koleo_weight = koleo_weight |
| 71 | |
| 72 | self.generator = torch.Generator() |
| 73 | self.step = -1 |
| 74 | |
| 75 | # Encoders |
| 76 | dino_head = partial(dino_head, in_dim=encoder.embed_dim) |
| 77 | if self.ibot_separate_head: |
| 78 | ibot_head = partial(dino_head, in_dim=encoder.embed_dim) |
| 79 | |
| 80 | self.student_encoder_dict, self.teacher_encoder_dict = dict(), dict() |
| 81 | self.student_encoder_dict["backbone"] = encoder |
| 82 | self.student_encoder_dict["dino_head"] = dino_head() |
| 83 | if self.ibot_separate_head: |
| 84 | self.student_encoder_dict['ibot_head'] = ibot_head() |
| 85 | self.student_encoder = nn.ModuleDict(self.student_encoder_dict) |
| 86 | |
| 87 | self.teacher_encoder_dict["backbone"] = copy.deepcopy(encoder) |
| 88 | self.teacher_encoder_dict["dino_head"] = dino_head() |
| 89 | if self.ibot_separate_head: |
nothing calls this directly
no test coverage detected