optimize fusion_form

Author: likyoo

change_detection_pytorch/base/decoder.py

@@ -2,6 +2,11 @@
 
 
 class Decoder(torch.nn.Module):
+    # TODO: support learnable fusion modules
+    def __init__(self):
+        super().__init__()
+        self.FUSION_DIC = {"2to1_fusion": ["sum", "diff", "abs_diff"],
+                           "2to2_fusion": ["concat"]}
 
     def fusion(self, x1, x2, fusion_form="concat"):
         """Specify the form of feature fusion"""
@@ -10,6 +15,8 @@ def fusion(self, x1, x2, fusion_form="concat"):
         elif fusion_form == "sum":
             x = x1 + x2
         elif fusion_form == "diff":
+            x = x2 - x1
+        elif fusion_form == "abs_diff":
             x = torch.abs(x1 - x2)
         else:
             raise ValueError('the fusion form "{}" is not defined'.format(fusion_form))

change_detection_pytorch/deeplabv3/decoder.py

@@ -42,7 +42,7 @@
 class DeepLabV3Decoder(nn.Sequential, Decoder):
     def __init__(self, in_channels, out_channels=256, atrous_rates=(12, 24, 36), fusion_form="concat"):
         # adjust encoder channels according to fusion form
-        if fusion_form == "concat":
+        if fusion_form in self.FUSION_DIC["2to2_fusion"]:
             in_channels = in_channels * 2
 
         super().__init__(
@@ -77,7 +77,7 @@ def __init__(
 
         # adjust encoder channels according to fusion form
         self.fusion_form = fusion_form
-        if self.fusion_form == "concat":
+        if self.fusion_form in self.FUSION_DIC["2to2_fusion"]:
             encoder_channels = [ch*2 for ch in encoder_channels]
 
         self.aspp = nn.Sequential(

change_detection_pytorch/fpn/decoder.py

@@ -97,7 +97,7 @@ def __init__(
 
         # adjust encoder channels according to fusion form
         self.fusion_form = fusion_form
-        if self.fusion_form == "concat":
+        if self.fusion_form in self.FUSION_DIC["2to2_fusion"]:
             encoder_channels = [ch*2 for ch in encoder_channels]
 
         self.p5 = nn.Conv2d(encoder_channels[0], pyramid_channels, kernel_size=1)

change_detection_pytorch/linknet/decoder.py

@@ -52,7 +52,7 @@ def __init__(
 
         # adjust encoder channels according to fusion form
         self.fusion_form = fusion_form
-        if self.fusion_form == "concat":
+        if self.fusion_form in self.FUSION_DIC["2to2_fusion"]:
             encoder_channels = [ch*2 for ch in encoder_channels]
 
         channels = list(encoder_channels) + [prefinal_channels]

change_detection_pytorch/manet/decoder.py

@@ -166,7 +166,7 @@ def __init__(
 
         # adjust encoder channels according to fusion form
         self.fusion_form = fusion_form
-        if self.fusion_form == "concat":
+        if self.fusion_form in self.FUSION_DIC["2to2_fusion"]:
             skip_channels = [ch*2 for ch in skip_channels]
             in_channels[0] = in_channels[0] * 2
             head_channels = head_channels * 2

change_detection_pytorch/pan/decoder.py

@@ -156,7 +156,7 @@ def __init__(
 
         # adjust encoder channels according to fusion form
         self.fusion_form = fusion_form
-        if self.fusion_form == "concat":
+        if self.fusion_form in self.FUSION_DIC["2to2_fusion"]:
             encoder_channels = [ch * 2 for ch in encoder_channels]
 
         self.fpa = FPABlock(in_channels=encoder_channels[-1], out_channels=decoder_channels)

change_detection_pytorch/pspnet/decoder.py

@@ -51,7 +51,7 @@ def __init__(
 
         # adjust encoder channels according to fusion form
         self.fusion_form = fusion_form
-        if self.fusion_form == "concat":
+        if self.fusion_form in self.FUSION_DIC["2to2_fusion"]:
             encoder_channels = [ch*2 for ch in encoder_channels]
 
         self.psp = PSPModule(

change_detection_pytorch/unet/decoder.py

@@ -96,7 +96,7 @@ def __init__(
 
         # adjust encoder channels according to fusion form
         self.fusion_form = fusion_form
-        if self.fusion_form == "concat":
+        if self.fusion_form in self.FUSION_DIC["2to2_fusion"]:
             skip_channels = [ch*2 for ch in skip_channels]
             in_channels[0] = in_channels[0] * 2
             head_channels = head_channels * 2

change_detection_pytorch/unetplusplus/decoder.py

@@ -94,7 +94,7 @@ def __init__(
 
         # adjust encoder channels according to fusion form
         self.fusion_form = fusion_form
-        if self.fusion_form == "concat":
+        if self.fusion_form in self.FUSION_DIC["2to2_fusion"]:
             self.skip_channels = [ch*2 for ch in self.skip_channels]
             self.in_channels[0] = self.in_channels[0] * 2
 

change_detection_pytorch/upernet/decoder.py

@@ -98,7 +98,7 @@ def __init__(
 
         # adjust encoder channels according to fusion form
         self.fusion_form = fusion_form
-        if self.fusion_form == "concat":
+        if self.fusion_form in self.FUSION_DIC["2to2_fusion"]:
             encoder_channels = [ch*2 for ch in encoder_channels]
 
         self.psp = PSPModule(