ResizeWith, UpsampleWith classes that can export to onnx with scale_factor in opset_v...
authorManu Mathew <a0393608@ti.com>
Tue, 3 Mar 2020 10:31:36 +0000 (16:01 +0530)
committerManu Mathew <a0393608@ti.com>
Tue, 3 Mar 2020 10:34:22 +0000 (16:04 +0530)
modules/pytorch_jacinto_ai/vision/models/pixel2pixel/deeplabv3lite.py
modules/pytorch_jacinto_ai/vision/models/pixel2pixel/fpnlite_pixel2pixel.py
modules/pytorch_jacinto_ai/vision/models/pixel2pixel/pixel2pixelnet_utils.py
modules/pytorch_jacinto_ai/vision/models/pixel2pixel/unetlite_pixel2pixel.py
modules/pytorch_jacinto_ai/xnn/layers/resize_blocks.py

index 5931e831ce63f3dfb85e9b79e2499e99cc58cd96..a7da59044b8ae665d69f9c2fc1d0347c7b896eba 100644 (file)
@@ -42,9 +42,9 @@ class DeepLabV3LiteDecoder(torch.nn.Module):
         self.decoder_channels = merged_channels = (current_channels+model_config.shortcut_out)
 
         upstride1 = model_config.shortcut_strides[-1]//model_config.shortcut_strides[0]
-        # use UpsampleScaleFactorGeneric() instead of UpsampleScaleFactor() to break down large upsampling factors to multiples of 4 and 2 -
+        # use UpsampleWithGeneric() instead of UpsampleWith() to break down large upsampling factors to multiples of 4 and 2 -
         # useful if upsampling factors other than 4 and 2 are not supported.
-        self.upsample1 = xnn.layers.UpsampleScaleFactor(decoder_channels, decoder_channels, upstride1, model_config.interpolation_type, model_config.interpolation_mode)
+        self.upsample1 = xnn.layers.UpsampleWith(decoder_channels, decoder_channels, upstride1, model_config.interpolation_type, model_config.interpolation_mode)
 
         self.cat = xnn.layers.CatBlock()
 
index 1ac77418c5ad991f15062d38e891f6001c33cc50..986b46643beeaaf9038a9d85f7b14f1638ca846f 100644 (file)
@@ -82,7 +82,7 @@ class FPNLitePyramid(torch.nn.Module):
             smooth_conv = xnn.layers.ConvDWSepNormAct2d(decoder_channels, decoder_channels, kernel_size=kernel_size_smooth, activation=(activation,activation)) \
                         if (inloop_fpn or all_outputs or is_last) else None
             self.smooth_convs.append(smooth_conv)
-            upsample = xnn.layers.UpsampleScaleFactor(decoder_channels, decoder_channels, upstride, interpolation_type, interpolation_mode)
+            upsample = xnn.layers.UpsampleWith(decoder_channels, decoder_channels, upstride, interpolation_type, interpolation_mode)
             self.upsamples.append(upsample)
         #
     #
index 972f708737930909305f7e7e7de221fe7dbce7fc..3073552365211a63424e89cae2d81e28a8701f57 100644 (file)
@@ -4,9 +4,9 @@ from .... import xnn
 def add_lite_prediction_modules(self, model_config, current_channels, module_names):
     # prediction and upsample
     if self.model_config.final_prediction:
-        # use UpsampleScaleFactorGeneric() instead of UpsampleScaleFactor(), to break down large upsampling factors to multiples of 4 and 2
+        # use UpsampleWithGeneric() instead of UpsampleWith(), to break down large upsampling factors to multiples of 4 and 2
         # useful if scale_factor other than 4 and 2 are not supported.
-        UpsampleClass = xnn.layers.UpsampleScaleFactor
+        UpsampleClass = xnn.layers.UpsampleWith
 
         # can control the range of final output with output_range
         final_activation = xnn.layers.get_fixed_pact2(output_range=model_config.output_range) if (model_config.output_range is not None) else False
index ed51cef35470ab7f803cc704305a5e7249862033..0d6c5b6da449d782fa00e7e724f2e939999ae099 100644 (file)
@@ -68,7 +68,7 @@ class UNetLitePyramid(torch.nn.Module):
         upstride = 2
         activation2 = (activation, activation)
         for idx, (s_stride, feat_chan) in enumerate(zip(shortcut_strides, shortcut_channels)):
-            self.upsamples.append(xnn.layers.UpsampleScaleFactor(current_channels, current_channels, upstride, interpolation_type, interpolation_mode))
+            self.upsamples.append(xnn.layers.UpsampleWith(current_channels, current_channels, upstride, interpolation_type, interpolation_mode))
             self.concats.append(xnn.layers.CatBlock())
             smooth_channels = max(minimum_channels, feat_chan)
             self.smooth_convs.append( xnn.layers.ConvDWSepNormAct2d(current_channels+feat_chan, smooth_channels, kernel_size=kernel_size_smooth, activation=activation2))
index ac0ee3ce556409a449047ffaf7cdbeb078abcf91..acd547af01c9444baa562a1315e697fe77b0bfc7 100644 (file)
@@ -1,3 +1,4 @@
+import warnings
 import torch
 from .deconv_blocks import *
 
@@ -11,7 +12,7 @@ from .deconv_blocks import *
 # onnx export from PyTorch is creating a complicated graph - use this workaround for now until the onnx export is fixed.
 # only way to create a simple graph with scale _factors seem to be provide size as integer to interpolate function
 # this workaround seems to be working in onnx opset_version=9, however in opset_version=11, it still produces a complicated graph.
-class ResizeScaleFactor(torch.nn.Module):
+class ResizeWith(torch.nn.Module):
     def __init__(self, scale_factor=None, mode='nearest'):
         ''' Resize with scale_factor
             This module exports an onnx graph with scale_factor
@@ -24,21 +25,28 @@ class ResizeScaleFactor(torch.nn.Module):
     def forward(self, x):
         assert isinstance(x, torch.Tensor), 'must provide a single tensor as input'
         scale_factor = (self.scale_factor, self.scale_factor) if not isinstance(self.scale_factor, (list,tuple)) else self.scale_factor
-        # generate size as a tuple and pass it - as onnx export inserts scale_factor if the size is a non-tensor
-        # this seems to be the only way to insert scale_factor in onnx export of Upsample/Resize
-        size = (int(x.shape[2]*scale_factor[0]), int(x.shape[3]*scale_factor[1]))
-        y = torch.nn.functional.interpolate(x, size=size, mode=self.mode)
+        try:
+            y = torch.nn.functional.interpolate(x, scale_factor=scale_factor, mode=self.mode, recompute_scale_factor = False)
+        except:
+            # warnings.warn('Note: If you are exporting to onnx_opset_version>=11, for a simple Upsample/Resize graph using scale_factor, please use pytorch>=1.5.' \
+            #               'Until pytorch 1.5 is available, you can install the nightly, as explained here: https://pytorch.org/blog/')
+            # The following trick seems to be the only way to export Upsample/Resize with scale_factor in onnx opset_version=9.
+            # Generate size as a tuple and pass it - as onnx export inserts scale_factor if the size is a non-tensor.
+            # This trick may not work in onnx opset_version=11, and we recommend to install pytorch 1.5 or latest nightly as explained in the warning above.
+            size = (int(x.shape[2]*scale_factor[0]), int(x.shape[3]*scale_factor[1]))
+            y = torch.nn.functional.interpolate(x, size=size, mode=self.mode)
+        #
         return y
 
 
-def UpsampleScaleFactor(input_channels=None, output_channels=None, upstride=None, interpolation_type='upsample', interpolation_mode='bilinear',
+def UpsampleWith(input_channels=None, output_channels=None, upstride=None, interpolation_type='upsample', interpolation_mode='bilinear',
                is_final_layer=False, final_activation=True):
     '''
          is_final_layer: Final layer in a pixel2pixel network should not typically use BatchNorm for best accuracy
          final_activation: use to control the range of final layer if required
      '''
     if interpolation_type == 'upsample':
-        upsample = ResizeScaleFactor(scale_factor=upstride, mode=interpolation_mode)
+        upsample = ResizeWith(scale_factor=upstride, mode=interpolation_mode)
     else:
         assert upstride is not None, 'upstride must not be None in this interpolation_mode'
         assert input_channels is not None, 'input_channels must not be None in this interpolation_mode'
@@ -50,7 +58,7 @@ def UpsampleScaleFactor(input_channels=None, output_channels=None, upstride=None
             upsample = [DeConvDWSepNormAct2d(input_channels, output_channels, kernel_size=upstride * 2, stride=upstride,
                                       normalization=normalization, activation=activation)]
         elif interpolation_type == 'upsample_conv':
-            upsample = [ResizeScaleFactor(scale_factor=upstride, mode=interpolation_mode),
+            upsample = [ResizeWith(scale_factor=upstride, mode=interpolation_mode),
                         ConvDWSepNormAct2d(input_channels, output_channels, kernel_size=int(upstride * 1.5 + 1),
                                       normalization=normalization, activation=activation)]
         elif interpolation_type == 'subpixel_conv':
@@ -66,7 +74,7 @@ def UpsampleScaleFactor(input_channels=None, output_channels=None, upstride=None
 
 
 
-class UpsampleScaleFactorGeneric(torch.nn.Module):
+class UpsampleWithGeneric(torch.nn.Module):
     def __init__(self, input_channels=None, output_channels=None, upstride=None, interpolation_type='upsample', interpolation_mode='bilinear',
                  is_final_layer=False, final_activation=False):
         '''
@@ -78,7 +86,7 @@ class UpsampleScaleFactorGeneric(torch.nn.Module):
         self.upsample_list = torch.nn.ModuleList()
         while upstride >= 2:
             upstride_layer = 4 if upstride > 4 else upstride
-            upsample = UpsampleScaleFactor(input_channels, output_channels, upstride_layer, interpolation_type, interpolation_mode,
+            upsample = UpsampleWith(input_channels, output_channels, upstride_layer, interpolation_type, interpolation_mode,
                                   is_final_layer=is_final_layer, final_activation=final_activation)
             self.upsample_list.append(upsample)
             upstride = upstride//4