update PAN decoder support encoder depth

Author: brianhou0208

segmentation_models_pytorch/decoders/pan/decoder.py

@@ -175,34 +175,40 @@ def forward(self, x, y):
 
 class PANDecoder(nn.Module):
     def __init__(
-        self, encoder_channels, decoder_channels, upscale_mode: str = "bilinear"
+        self,
+        encoder_channels,
+        encoder_depth,
+        decoder_channels,
+        upscale_mode: str = "bilinear",
     ):
         super().__init__()
 
+        if encoder_depth < 3:
+            raise ValueError(
+                "Encoder depth for PAN decoder cannot be less than 3, got {}.".format(
+                    encoder_depth
+                )
+            )
+
+        encoder_channels = encoder_channels[2:][::-1]
+
         self.fpa = FPABlock(
-            in_channels=encoder_channels[-1], out_channels=decoder_channels
-        )
-        self.gau3 = GAUBlock(
-            in_channels=encoder_channels[-2],
-            out_channels=decoder_channels,
-            upscale_mode=upscale_mode,
-        )
-        self.gau2 = GAUBlock(
-            in_channels=encoder_channels[-3],
-            out_channels=decoder_channels,
-            upscale_mode=upscale_mode,
-        )
-        self.gau1 = GAUBlock(
-            in_channels=encoder_channels[-4],
-            out_channels=decoder_channels,
-            upscale_mode=upscale_mode,
+            in_channels=encoder_channels[0], out_channels=decoder_channels
         )
 
+        for i in range(1, len(encoder_channels)):
+            self.add_module(f"gau{len(encoder_channels)-i}", GAUBlock(
+                in_channels=encoder_channels[i],
+                out_channels=decoder_channels,
+                upscale_mode=upscale_mode,
+            ))
+
     def forward(self, *features):
-        bottleneck = features[-1]
-        x5 = self.fpa(bottleneck)  # 1/32
-        x4 = self.gau3(features[-2], x5)  # 1/16
-        x3 = self.gau2(features[-3], x4)  # 1/8
-        x2 = self.gau1(features[-4], x3)  # 1/4
+        features = features[2:]  # remove first and second skip
+        features = features[::-1]  # reverse channels to start from head of encoder
+
+        out = self.fpa(features[0])
 
-        return x2
+        for i in range(1, len(features)):
+            out = getattr(self, f"gau{len(features)-i}")(features[i], out)
+        return out