use AddBlock for addition in fpn
authorManu Mathew <a0393608@ti.com>
Mon, 14 Sep 2020 07:49:56 +0000 (13:19 +0530)
committerManu Mathew <a0393608@ti.com>
Tue, 15 Sep 2020 06:46:11 +0000 (12:16 +0530)
modules/pytorch_jacinto_ai/xvision/models/pixel2pixel/fpnlite_pixel2pixel.py

index 293325262ce25dbdb4754e9052648183d136ed99..9427f78788e0f470e1a822727b9d52f22bf8169b 100644 (file)
@@ -70,6 +70,7 @@ class FPNLitePyramid(torch.nn.Module):
         self.smooth_convs = torch.nn.ModuleList()
         self.shortcuts = torch.nn.ModuleList()
         self.upsamples = torch.nn.ModuleList()
+        self.adds = torch.nn.ModuleList()
 
         shortcut0 = self.create_shortcut(current_channels, decoder_channels, activation) \
             if (current_channels != decoder_channels) else None
@@ -94,6 +95,7 @@ class FPNLitePyramid(torch.nn.Module):
             upsample = xnn.layers.UpsampleWith(decoder_channels, decoder_channels, upstride, interpolation_type,
                                                interpolation_mode)
             self.upsamples.append(upsample)
+            self.adds.append(xnn.layers.AddBlock())
         #
     #
 
@@ -121,7 +123,7 @@ class FPNLitePyramid(torch.nn.Module):
             x_s = shortcut(x_s)
             # updample current output and add to that
             x = upsample(x)
-            x = x + x_s
+            x = self.adds[idx]((x,x_s))
             # smooth conv
             y = smooth_conv(x) if (smooth_conv is not None) else x
             # use smooth output for next level in inloop_fpn