Update HRNet to add first_conv_stride_one parameter in all variants

Author: BloodAxe

pytorch_toolbelt/modules/encoders/timm/hrnet.py

@@ -5,25 +5,14 @@
 __all__ = ["HRNetW18Encoder", "HRNetW32Encoder", "HRNetW48Encoder", "TimmHRNetW18SmallV2Encoder"]
 
 
-class HRNetW18Encoder(GenericTimmEncoder):
-    def __init__(
-        self, pretrained=True, use_incre_features: bool = True, layers=None, first_conv_stride_one: bool = False
-    ):
-        from timm.models import hrnet
-
-        encoder = hrnet.hrnet_w18(
-            pretrained=pretrained,
-            feature_location="incre" if use_incre_features else "",
-            features_only=True,
-            out_indices=(0, 1, 2, 3, 4),
-        )
+class HRNetTimmEncoder(GenericTimmEncoder):
+    def __init__(self, encoder, first_conv_stride_one, layers):
         if first_conv_stride_one:
             encoder.conv1.stride = (1, 1)
 
         super().__init__(encoder, layers)
-
         if first_conv_stride_one:
-            self._output_strides = [s // 2 for s in self._output_strides]
+            self._output_strides = (s // 2 for s in self._output_strides)
 
     def forward(self, x):
         y = self.encoder.forward(x)
@@ -34,29 +23,40 @@ def change_input_channels(self, input_channels: int, mode="auto", **kwargs):
         return self
 
 
-class HRNetW32Encoder(GenericTimmEncoder):
-    def __init__(self, pretrained=True, use_incre_features: bool = True, layers=None):
+class HRNetW18Encoder(HRNetTimmEncoder):
+    def __init__(
+        self, pretrained=True, use_incre_features: bool = True, layers=None, first_conv_stride_one: bool = False
+    ):
         from timm.models import hrnet
 
-        encoder = hrnet.hrnet_w32(
+        encoder = hrnet.hrnet_w18(
             pretrained=pretrained,
             feature_location="incre" if use_incre_features else "",
             features_only=True,
             out_indices=(0, 1, 2, 3, 4),
         )
-        super().__init__(encoder, layers)
+        super().__init__(encoder, first_conv_stride_one=first_conv_stride_one, layers=layers)
 
-    def forward(self, x):
-        y = self.encoder.forward(x)
-        return _take(y, self._layers)
 
-    def change_input_channels(self, input_channels: int, mode="auto", **kwargs):
-        self.encoder.conv1 = make_n_channel_input(self.encoder.conv1, input_channels, mode, **kwargs)
-        return self
+class HRNetW32Encoder(HRNetTimmEncoder):
+    def __init__(
+        elf, pretrained=True, use_incre_features: bool = True, layers=None, first_conv_stride_one: bool = False
+    ):
+        from timm.models import hrnet
 
+        encoder = hrnet.hrnet_w32(
+            pretrained=pretrained,
+            feature_location="incre" if use_incre_features else "",
+            features_only=True,
+            out_indices=(0, 1, 2, 3, 4),
+        )
+        super().__init__(encoder, first_conv_stride_one=first_conv_stride_one, layers=layers)
 
-class HRNetW48Encoder(GenericTimmEncoder):
-    def __init__(self, pretrained=True, use_incre_features: bool = True, layers=None):
+
+class HRNetW48Encoder(HRNetTimmEncoder):
+    def __init__(
+        elf, pretrained=True, use_incre_features: bool = True, layers=None, first_conv_stride_one: bool = False
+    ):
         from timm.models import hrnet
 
         encoder = hrnet.hrnet_w48(
@@ -65,19 +65,13 @@ def __init__(self, pretrained=True, use_incre_features: bool = True, layers=None
             features_only=True,
             out_indices=(0, 1, 2, 3, 4),
         )
-        super().__init__(encoder, layers)
-
-    def forward(self, x):
-        y = self.encoder.forward(x)
-        return _take(y, self._layers)
-
-    def change_input_channels(self, input_channels: int, mode="auto", **kwargs):
-        self.encoder.conv1 = make_n_channel_input(self.encoder.conv1, input_channels, mode, **kwargs)
-        return self
+        super().__init__(encoder, first_conv_stride_one=first_conv_stride_one, layers=layers)
 
 
-class TimmHRNetW18SmallV2Encoder(GenericTimmEncoder):
-    def __init__(self, pretrained=True, use_incre_features: bool = True, layers=None, activation=ACT_RELU):
+class TimmHRNetW18SmallV2Encoder(HRNetTimmEncoder):
+    def __init__(
+        self, elf, pretrained=True, use_incre_features: bool = True, layers=None, first_conv_stride_one: bool = False
+    ):
         from timm.models import hrnet
 
         encoder = hrnet.hrnet_w18_small_v2(
@@ -86,8 +80,4 @@ def __init__(self, pretrained=True, use_incre_features: bool = True, layers=None
             features_only=True,
             out_indices=(0, 1, 2, 3, 4),
         )
-        super().__init__(encoder, layers)
-
-    def change_input_channels(self, input_channels: int, mode="auto", **kwargs):
-        self.encoder.conv1 = make_n_channel_input(self.encoder.conv1, input_channels, mode, **kwargs)
-        return self
+        super().__init__(encoder, first_conv_stride_one=first_conv_stride_one, layers=layers)