Add encoder_kwargs to DeepLab and Linknet

Author: dpieczynski

segmentation_models_pytorch/decoders/deeplabv3/model.py

@@ -56,6 +56,7 @@ def __init__(
         activation: Optional[str] = None,
         upsampling: int = 8,
         aux_params: Optional[dict] = None,
+        encoder_kwargs: dict = {}
     ):
         super().__init__()
 
@@ -65,6 +66,7 @@ def __init__(
             depth=encoder_depth,
             weights=encoder_weights,
             output_stride=8,
+            **encoder_kwargs
         )
 
         self.decoder = DeepLabV3Decoder(
@@ -137,6 +139,7 @@ def __init__(
         activation: Optional[str] = None,
         upsampling: int = 4,
         aux_params: Optional[dict] = None,
+        encoder_kwargs: dict = {}
     ):
         super().__init__()
 
@@ -149,6 +152,7 @@ def __init__(
             depth=encoder_depth,
             weights=encoder_weights,
             output_stride=encoder_output_stride,
+            **encoder_kwargs
         )
 
         self.decoder = DeepLabV3PlusDecoder(

segmentation_models_pytorch/decoders/linknet/model.py

@@ -61,6 +61,7 @@ def __init__(
         classes: int = 1,
         activation: Optional[Union[str, callable]] = None,
         aux_params: Optional[dict] = None,
+        encoder_kwargs: dict = {}
     ):
         super().__init__()
 
@@ -69,6 +70,7 @@ def __init__(
             in_channels=in_channels,
             depth=encoder_depth,
             weights=encoder_weights,
+            **encoder_kwargs
         )
 
         self.decoder = LinknetDecoder(

segmentation_models_pytorch/encoders/timm_universal.py

@@ -3,9 +3,10 @@
 
 
 class TimmUniversalEncoder(nn.Module):
-    def __init__(self, name, pretrained=True, in_channels=3, depth=5, output_stride=32):
+    def __init__(self, name, pretrained=True, in_channels=3, depth=5, output_stride=32, **kwargs):
         super().__init__()
         kwargs = dict(
+            **kwargs,
             in_chans=in_channels,
             features_only=True,
             output_stride=output_stride,