PyTorch Lightning implementation of `Moco `_ Paper authors: Xinlei Chen, Haoqi Fan, Ross Girshick, Kaiming He. Code adapted from `facebookresearch/moco `_ to Lightning by: - `William Falcon <https://github.
| 10 | |
| 11 | |
| 12 | class Moco_v2(L.LightningModule): |
| 13 | """PyTorch Lightning implementation of `Moco <https://arxiv.org/abs/2003.04297>`_ |
| 14 | |
| 15 | Paper authors: Xinlei Chen, Haoqi Fan, Ross Girshick, Kaiming He. |
| 16 | |
| 17 | Code adapted from `facebookresearch/moco <https://github.com/facebookresearch/moco>`_ to Lightning by: |
| 18 | - `William Falcon <https://github.com/williamFalcon>`_ |
| 19 | """ |
| 20 | |
| 21 | def __init__( |
| 22 | self, |
| 23 | base_encoder: Union[str, torch.nn.Module] = "resnet18", |
| 24 | emb_dim: int = 128, |
| 25 | num_negatives: int = 65536, |
| 26 | encoder_momentum: float = 0.999, |
| 27 | softmax_temperature: float = 0.07, |
| 28 | learning_rate: float = 0.03, |
| 29 | momentum: float = 0.9, |
| 30 | weight_decay: float = 1e-4, |
| 31 | data_dir: str = "./", |
| 32 | batch_size: int = 256, |
| 33 | use_mlp: bool = False, |
| 34 | num_workers: int = 8, |
| 35 | *args, |
| 36 | **kwargs |
| 37 | ): |
| 38 | super().__init__() |
| 39 | self.save_hyperparameters() |
| 40 | |
| 41 | # create the encoders |
| 42 | # num_classes is the output fc dimension |
| 43 | self.encoder_q, self.encoder_k = self.init_encoders(base_encoder) |
| 44 | |
| 45 | if use_mlp: # hack: brute-force replacement |
| 46 | dim_mlp = self.encoder_q.fc.weight.shape[1] |
| 47 | self.encoder_q.fc = nn.Sequential( |
| 48 | nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_q.fc |
| 49 | ) |
| 50 | self.encoder_k.fc = nn.Sequential( |
| 51 | nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_k.fc |
| 52 | ) |
| 53 | |
| 54 | for param_q, param_k in zip( |
| 55 | self.encoder_q.parameters(), self.encoder_k.parameters() |
| 56 | ): |
| 57 | param_k.data.copy_(param_q.data) # initialize |
| 58 | param_k.requires_grad = False # not update by gradient |
| 59 | |
| 60 | # create the queue |
| 61 | self.register_buffer("queue", torch.randn(emb_dim, num_negatives)) |
| 62 | self.queue = nn.functional.normalize(self.queue, dim=0) |
| 63 | |
| 64 | self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long)) |
| 65 | |
| 66 | # create the validation queue |
| 67 | self.register_buffer("val_queue", torch.randn(emb_dim, num_negatives)) |
| 68 | self.val_queue = nn.functional.normalize(self.val_queue, dim=0) |
| 69 |
nothing calls this directly
no outgoing calls
no test coverage detected