Refactor model subclassing

Author: qubvel

segmentation_models_pytorch/unet/model.py

@@ -1,8 +1,5 @@
-import torch
-import torch.nn as nn
 from .decoder import UnetDecoder
-
-from ..base.model import EncoderDecoder
+from ..base import EncoderDecoder
 from ..encoders import get_encoder
 
 
@@ -18,7 +15,6 @@ def __init__(
             activation='sigmoid',
             center=False,  # usefull for VGG models
     ):
-
         encoder = get_encoder(
             encoder_name,
             encoder_weights=encoder_weights
@@ -32,23 +28,6 @@ def __init__(
             center=center,
         )
 
-        # define activation function
-        if activation == 'softmax':
-            activation_fn = nn.Softmax(dim=1)
-        elif activation == 'sigmoid':
-            activation_fn = nn.Sigmoid()
-        else:
-            raise ValueError('Activation should be "sigmoid" or "softmax"')
-
-        super().__init__(encoder, decoder, activation_fn)
-
-        if encoder_weights is not None:
-            self.set_preprocessing_params(
-                input_size=encoder.input_size,
-                input_space=encoder.input_space,
-                input_range=encoder.input_range,
-                mean=encoder.mean,
-                std=encoder.std,
-            )
+        super().__init__(encoder, decoder, activation)
 
         self.name = 'u-{}'.format(encoder_name)