fix docstring

openselfsup/models/byol.py

@@ -15,9 +15,10 @@ class BYOL(nn.Module):
     Self-Supervised Learning (https://arxiv.org/abs/2006.07733)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
         base_momentum (float): The base momentum coefficient for the target network.
             Default: 0.996.

openselfsup/models/classification.py

@@ -14,9 +14,9 @@ class Classification(nn.Module):
     """Simple image classification.
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
+        backbone (dict): Config dict for module of backbone ConvNet.
         with_sobel (bool): Whether to apply a Sobel filter on images. Default: False.
-        head (nn.Module): Module of loss functions.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 

openselfsup/models/deepcluster.py

@@ -18,10 +18,11 @@ class DeepCluster(nn.Module):
     of Visual Features (https://arxiv.org/abs/1807.05520)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
+        backbone (dict): Config dict for module of backbone ConvNet.
         with_sobel (bool): Whether to apply a Sobel filter on images. Default: False.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 

openselfsup/models/moco.py

@@ -17,9 +17,10 @@ class MOCO(nn.Module):
     "https://github.com/facebookresearch/moco/blob/master/moco/builder.py".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
         queue_len (int): Number of negative keys maintained in the queue.
             Default: 65536.

openselfsup/models/npid.py

@@ -15,10 +15,11 @@ class NPID(nn.Module):
     Instance Discrimination (https://arxiv.org/abs/1805.01978)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
-        memory_bank (nn.Module): Module of memory banks.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
+        memory_bank (dict): Config dict for module of memory banks. Default: None.
         neg_num (int): Number of negative samples for each image. Default: 65536.
         ensure_neg (bool): If False, there is a small probability
             that negative samples contain positive ones. Default: False.

openselfsup/models/odc.py

@@ -17,11 +17,12 @@ class ODC(nn.Module):
     (https://arxiv.org/abs/2006.10645)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
+        backbone (dict): Config dict for module of backbone ConvNet.
         with_sobel (bool): Whether to apply a Sobel filter on images. Default: False.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
-        memory_bank (nn.Module): Module of memory banks.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
+        memory_bank (dict): Module of memory banks. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 

openselfsup/models/relative_loc.py

@@ -15,9 +15,10 @@ class RelativeLoc(nn.Module):
     by Context Prediction (https://arxiv.org/abs/1505.05192)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 

openselfsup/models/rotation_pred.py

@@ -15,8 +15,8 @@ class RotationPred(nn.Module):
     by Predicting Image Rotations (https://arxiv.org/abs/1803.07728)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 

openselfsup/models/simclr.py

@@ -16,9 +16,10 @@ class SimCLR(nn.Module):
     of Visual Representations (https://arxiv.org/abs/2002.05709)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """