Docstrings updated

Author: qubvel

segmentation_models_pytorch/base/encoder_decoder.py

@@ -9,22 +9,33 @@ def __init__(self, encoder, decoder, activation):
         super().__init__()
         self.encoder = encoder
         self.decoder = decoder
-        
+
         if callable(activation) or activation is None:
             self.activation = activation
         elif activation == 'softmax':
             self.activation = nn.Softmax(dim=1)
         elif activation == 'sigmoid':
             self.activation = nn.Sigmoid()
         else:
-            raise ValueError('Activation should be "sigmoid" or "softmax"')
-            
+            raise ValueError('Activation should be "sigmoid"/"softmax"/callable/None')
+
     def forward(self, x):
+        """Sequentially pass `x` trough model`s `encoder` and `decoder` (return logits!)"""
         x = self.encoder(x)
         x = self.decoder(x)
         return x
 
     def predict(self, x):
+        """Inference method. Switch model to `eval` mode, call `.forward(x)`
+        and apply activation function (if activation is not `None`) with `torch.no_grad()`
+
+        Args:
+            x: 4D torch tensor with shape (batch_size, channels, height, width)
+
+        Return:
+            prediction: 4D torch tensor with shape (batch_size, classes, height, width)
+
+        """
         if self.training:
             self.eval()
 

segmentation_models_pytorch/fpn/model.py

@@ -13,10 +13,11 @@ class FPN(EncoderDecoder):
         decoder_segmentation_channels: a number of convolution filters in segmentation head of FPN_.
         classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
         dropout: spatial dropout rate in range (0, 1).
-        activation: one of [``sigmoid``, ``softmax``, None]
+        activation: activation function used in ``.predict(x)`` method for inference.
+            One of [``sigmoid``, ``softmax``, callable, None]
 
     Returns:
-        ``keras.models.Model``: **FPN**
+        ``torch.nn.Module``: **FPN**
 
     .. _FPN:
         http://presentations.cocodataset.org/COCO17-Stuff-FAIR.pdf

segmentation_models_pytorch/linknet/model.py

@@ -16,7 +16,8 @@ class Linknet(EncoderDecoder):
         decoder_use_batchnorm: if ``True``, ``BatchNormalisation`` layer between ``Conv2D`` and ``Activation`` layers
             is used.
         classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
-        activation: one of [``sigmoid``, ``softmax``, None]
+        activation: activation function used in ``.predict(x)`` method for inference.
+            One of [``sigmoid``, ``softmax``, callable, None]
 
     Returns:
         ``torch.nn.Module``: **Linknet**

segmentation_models_pytorch/pspnet/model.py

@@ -18,7 +18,8 @@ class PSPNet(EncoderDecoder):
         psp_aux_output: if ``True`` add auxiliary classification output for encoder training
         psp_dropout: spatial dropout rate between 0 and 1.
         classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
-        activation: one of [``sigmoid``, ``softmax``, None]
+        activation: activation function used in ``.predict(x)`` method for inference.
+            One of [``sigmoid``, ``softmax``, callable, None]
 
     Returns:
         ``torch.nn.Module``: **PSPNet**

segmentation_models_pytorch/unet/model.py

@@ -14,7 +14,8 @@ class Unet(EncoderDecoder):
         decoder_use_batchnorm: if ``True``, ``BatchNormalisation`` layer between ``Conv2D`` and ``Activation`` layers
             is used.
         classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
-        activation: one of [``sigmoid``, ``softmax``, None]
+        activation: activation function used in ``.predict(x)`` method for inference.
+            One of [``sigmoid``, ``softmax``, callable, None]
         center: if ``True`` add ``Conv2dReLU`` block on encoder head (useful for VGG models)
 
     Returns: