fix minor issues with NNBYOL and NNSiam

README.md

@@ -15,7 +15,7 @@ While the library is self contained, it is possible to use the models outside of
 ---
 
 ## News
-* **[Sep 27 2021]**: :pizza: Added [NNSiam](https://arxiv.org/abs/2104.14548), NNBYOL, new tutorials for implementing new methods [1](https://solo-learn.readthedocs.io/en/latest/tutorials/add_new_method.html) and [2](https://solo-learn.readthedocs.io/en/latest/tutorials/add_new_method_momentum.html), more testing and fixed issues with custom data and linear evaluation.
+* **[Sep 27 2021]**: :pizza: Added [NNSiam](https://arxiv.org/abs/2104.14548), [NNBYOL]((https://arxiv.org/abs/2104.14548), new tutorials for implementing new methods [1](https://solo-learn.readthedocs.io/en/latest/tutorials/add_new_method.html) and [2](https://solo-learn.readthedocs.io/en/latest/tutorials/add_new_method_momentum.html), more testing and fixed issues with custom data and linear evaluation.
 * **[Sep 19 2021]**: :kangaroo: Added online k-NN evaluation.
 * **[Sep 17 2021]**: :robot: Added [ViT](https://arxiv.org/abs/2010.11929) and [Swin](https://arxiv.org/abs/2103.14030).
 * **[Sep 13 2021]**: :book: Improved [Docs](https://solo-learn.readthedocs.io/en/latest/?badge=latest) and added tutorials for [pretraining](https://solo-learn.readthedocs.io/en/latest/tutorials/overview.html) and [offline linear eval](https://solo-learn.readthedocs.io/en/latest/tutorials/offline_linear_eval.html).
@@ -33,7 +33,9 @@ While the library is self contained, it is possible to use the models outside of
 * [DeepCluster V2](https://arxiv.org/abs/2006.09882)
 * [DINO](https://arxiv.org/abs/2104.14294)
 * [MoCo V2+](https://arxiv.org/abs/2003.04297)
+* [NNBYOL](https://arxiv.org/abs/2104.14548)
 * [NNCLR](https://arxiv.org/abs/2104.14548)
+* [NNSiam](https://arxiv.org/abs/2104.14548)
 * [ReSSL](https://arxiv.org/abs/2107.09282)
 * [SimCLR](https://arxiv.org/abs/2002.05709) + [Supervised Contrastive Learning](https://arxiv.org/abs/2004.11362)
 * [SimSiam](https://arxiv.org/abs/2011.10566)

docs/source/index.rst

@@ -69,7 +69,9 @@ While the library is self contained, it is possible to use the models outside of
    solo/methods/deepclusterv2
    solo/methods/dino
    solo/methods/mocov2plus
+   solo/methods/nnbyol
    solo/methods/nnclr
+   solo/methods/nnsiam
    solo/methods/ressl
    solo/methods/simclr
    solo/methods/simsiam

docs/source/solo/methods/nnbyol.rst

@@ -0,0 +1,36 @@
+NNBYOL
+======
+
+.. automethod:: solo.methods.nnbyol.NNBYOL.__init__
+   :noindex:
+
+
+add_model_specific_args
+~~~~~~~~~~~~~~~~~~~~~~~
+.. automethod:: solo.methods.nnbyol.NNBYOL.add_model_specific_args
+   :noindex:
+
+learnable_params
+~~~~~~~~~~~~~~~~
+.. autoattribute:: solo.methods.nnbyol.NNBYOL.learnable_params
+   :noindex:
+
+dequeue_and_enqueue
+~~~~~~~~~~~~~~~~~~~
+.. automethod:: solo.methods.nnbyol.NNBYOL.dequeue_and_enqueue
+   :noindex:
+
+find_nn
+~~~~~~~~~~~~~~~~~~~~
+.. automethod:: solo.methods.nnbyol.NNBYOL.find_nn
+   :noindex:
+
+forward
+~~~~~~~
+.. automethod:: solo.methods.nnbyol.NNBYOL.forward
+   :noindex:
+
+training_step
+~~~~~~~~~~~~~
+.. automethod:: solo.methods.nnbyol.NNBYOL.training_step
+   :noindex:

docs/source/solo/methods/nnsiam.rst

@@ -0,0 +1,36 @@
+NNSiam
+======
+
+.. automethod:: solo.methods.nnsiam.NNSiam.__init__
+   :noindex:
+
+
+add_model_specific_args
+~~~~~~~~~~~~~~~~~~~~~~~
+.. automethod:: solo.methods.nnsiam.NNSiam.add_model_specific_args
+   :noindex:
+
+learnable_params
+~~~~~~~~~~~~~~~~
+.. autoattribute:: solo.methods.nnsiam.NNSiam.learnable_params
+   :noindex:
+
+dequeue_and_enqueue
+~~~~~~~~~~~~~~~~~~~
+.. automethod:: solo.methods.nnsiam.NNSiam.dequeue_and_enqueue
+   :noindex:
+
+find_nn
+~~~~~~~~~~~~~~~~~~~~
+.. automethod:: solo.methods.nnsiam.NNSiam.find_nn
+   :noindex:
+
+forward
+~~~~~~~
+.. automethod:: solo.methods.nnsiam.NNSiam.forward
+   :noindex:
+
+training_step
+~~~~~~~~~~~~~
+.. automethod:: solo.methods.nnsiam.NNSiam.training_step
+   :noindex:

solo/methods/nnbyol.py

@@ -45,6 +45,12 @@ def __init__(
             proj_hidden_dim (int): number of neurons of the hidden layers of the projector.
             pred_hidden_dim (int): number of neurons of the hidden layers of the predictor.
             queue_size (int): number of samples to keep in the queue.
+
+        .. note::
+            NNBYOL is similar to NNSiam but the queue from which the neighbors are retrieved is
+            updated using the features of the momentum encoder. See NNCLR's paper for more detals:
+            https://arxiv.org/abs/2104.14548
+
         """
 
         super().__init__(**kwargs)
@@ -180,15 +186,15 @@ def forward(self, X: torch.Tensor, *args, **kwargs) -> Dict[str, Any]:
         return {**out, "z": z, "p": p}
 
     def training_step(self, batch: Sequence[Any], batch_idx: int) -> torch.Tensor:
-        """Training step for BYOL reusing BaseMethod training step.
+        """Training step for NNBYOL reusing BaseMethod training step.
 
         Args:
             batch (Sequence[Any]): a batch of data in the format of [img_indexes, [X], Y], where
                 [X] is a list of size self.num_crops containing batches of images.
             batch_idx (int): index of the batch.
 
         Returns:
-            torch.Tensor: total loss composed of BYOL and classification loss.
+            torch.Tensor: total loss composed of NNBYOL and classification loss.
         """
 
         targets = batch[-1]

solo/methods/nnsiam.py

@@ -174,7 +174,7 @@ def training_step(self, batch: Sequence[Any], batch_idx: int) -> torch.Tensor:
             batch_idx (int): index of the batch
 
         Returns:
-            torch.Tensor: total loss composed of SimSiam loss and classification loss
+            torch.Tensor: total loss composed of NNSiam loss and classification loss
         """
 
         targets = batch[-1]