🐛 [Bug] Cannot convert simple torchscript containing two torch.nn.Upsample operations

[gcuendet]

Bug Description

Scripting a simple "network" containing two torch.nn.Upsample modules and trying to convert the resulting torchscript does not work.

To Reproduce

Steps to reproduce the behavior:

1. Generate a torchscript of the Network below, with torch.jit.script.
2. Try to convert to TensorRT with torch_tensorrt (I tried both in C++, on linux and in python)

class Network(torch.nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        # NOTE:
        # * Specifying a single float as scale_factor or the 2d tuple doesn't change the behavior
        self.upsample1 = torch.nn.Upsample(
            scale_factor=(2.0, 2.0), mode="bilinear", align_corners=False
        )
        self.upsample2 = torch.nn.Upsample(
            scale_factor=(2.0, 2.0), mode="bilinear", align_corners=False
        )

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        # NOTE:
        # * Using the same self.upsample() module doesn't work as well
        # * Computing and returning out doesn't work as well
        out1 = self.upsample1(X)
        out2 = self.upsample2(X)
        # out = out1 + out2
        return out1

Expected behavior

The conversions succeeds and a new valid torchscript is obtained.

Environment

I managed to reproduce the bug both when using pytorch 1.11 and torch-tensorRT 1.1.0 and using pytorch 1.13.1 and torch-tensorRT main.

Torch-TensorRT 1.1.0

• Torch-TensorRT Version (e.g. 1.0.0): 1.1.0
• PyTorch Version (e.g. 1.0): 1.11
• CPU Architecture: x86_64
• OS (e.g., Linux): Linux, Ubuntu 20.04
• How you installed PyTorch (conda, pip, libtorch, source): pip for the python package used to generate the torchscript, source for the C++ dependency linked to Torch-TensorRT
• Build command you used (if compiling from source): Compiled with CMake
• Python version: 3.8.16
• CUDA version: 11.3
• GPU models and configuration: NVIDIA GeForce RTX 2070 SUPER, SM Capability: 7.5
• Any other relevant information:

When using torch-tensorRT 1.1.0, I get the following error:

DEBUG: [Torch-TensorRT] - Registering input/output torch::jit::Value for segmented graphs
terminate called after throwing an instance of 'c10::Error'
  what():  Expected Tensor but got Uninitialized


Exception raised from reportToTensorTypeError at /home/ubuntu/buildAgent/temp/buildTmp/conan_home/.conan/data/libtorch/1.11.0-5/cognex/stable/build/b8fcba865a68ae644e8a8bfa868ef00d13dc8a17/source_subfolder/aten/src/ATen/core/ivalue.cpp:908 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0x6b (0x7f66313e708b in /mnt/caches/conan/data/libtorch/1.11.0-5/cognex/stable/package/b8fcba865a68ae644e8a8bfa868ef00d13dc8a17/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0xce (0x7f66313e2a5e in /mnt/caches/conan/data/libtorch/1.11.0-5/cognex/stable/package/b8fcba865a68ae644e8a8bfa868ef00d13dc8a17/lib/libc10.so)
frame #2: c10::IValue::reportToTensorTypeError() const + 0x64 (0x7f6634480064 in /mnt/caches/conan/data/libtorch/1.11.0-5/cognex/stable/package/b8fcba865a68ae644e8a8bfa868ef00d13dc8a17/lib/libtorch_cpu.so)
frame #3: torch_tensorrt::core::partitioning::getSegmentsOutputByRunning(torch_tensorrt::core::partitioning::SegmentedBlock&, std::unordered_map<torch::jit::Value const*, c10::IValue, std::hash<torch::jit::Value const*>, std::equal_to<torch::jit::Value const*>, std::allocator<std::pair<torch::jit::Value const* const, c10::IValue> > >&, torch_tensorrt::core::partitioning::PartitionInfo const&) + 0x15a7 (0x7f6639fd8797 in /mnt/caches/conan/data/torch-tensorrt/1.1.0-6/cognex/stable/package/c2bc9df050867ddb44d73bc998d3c8ceef4f763a/lib/libtorchtrt.so)
frame #4: torch_tensorrt::core::partitioning::runShapeAnalysis(std::vector<torch_tensorrt::core::partitioning::SegmentedBlock, std::allocator<torch_tensorrt::core::partitioning::SegmentedBlock> >&, std::unordered_map<torch::jit::Value const*, c10::IValue, std::hash<torch::jit::Value const*>, std::equal_to<torch::jit::Value const*>, std::allocator<std::pair<torch::jit::Value const* const, c10::IValue> > >&, torch_tensorrt::core::partitioning::PartitionInfo const&) + 0x81 (0x7f6639fda581 in /mnt/caches/conan/data/torch-tensorrt/1.1.0-6/cognex/stable/package/c2bc9df050867ddb44d73bc998d3c8ceef4f763a/lib/libtorchtrt.so)
frame #5: torch_tensorrt::core::partitioning::Partition(torch::jit::Block*, std::unordered_map<torch::jit::Value const*, c10::IValue, std::hash<torch::jit::Value const*>, std::equal_to<torch::jit::Value const*>, std::allocator<std::pair<torch::jit::Value const* const, c10::IValue> > >&, torch_tensorrt::core::partitioning::PartitionInfo const&) + 0x19b (0x7f6639fe587b in /mnt/caches/conan/data/torch-tensorrt/1.1.0-6/cognex/stable/package/c2bc9df050867ddb44d73bc998d3c8ceef4f763a/lib/libtorchtrt.so)
frame #6: torch_tensorrt::core::ConstructFallbackGraph(torch::jit::Module&, torch::jit::Block*, std::unordered_map<torch::jit::Value const*, c10::IValue, std::hash<torch::jit::Value const*>, std::equal_to<torch::jit::Value const*>, std::allocator<std::pair<torch::jit::Value const* const, c10::IValue> > >, torch_tensorrt::core::CompileSpec, std::map<torch::jit::Value*, c10::IValue, std::less<torch::jit::Value*>, std::allocator<std::pair<torch::jit::Value* const, c10::IValue> > >) + 0xff (0x7f663a00006f in /mnt/caches/conan/data/torch-tensorrt/1.1.0-6/cognex/stable/package/c2bc9df050867ddb44d73bc998d3c8ceef4f763a/lib/libtorchtrt.so)
frame #7: torch_tensorrt::core::CompileGraph(torch::jit::Module const&, torch_tensorrt::core::CompileSpec) + 0x980 (0x7f663a002a00 in /mnt/caches/conan/data/torch-tensorrt/1.1.0-6/cognex/stable/package/c2bc9df050867ddb44d73bc998d3c8ceef4f763a/lib/libtorchtrt.so)
frame #8: torch_tensorrt::torchscript::compile(torch::jit::Module const&, torch_tensorrt::torchscript::CompileSpec) + 0x5b7 (0x7f6639e9e7f7 in /mnt/caches/conan/data/torch-tensorrt/1.1.0-6/cognex/stable/package/c2bc9df050867ddb44d73bc998d3c8ceef4f763a/lib/libtorchtrt.so)
frame #9: <unknown function> + 0x5ec8 (0x56183fd18ec8 in ./build/src/interpolate_tensorrt)
frame #10: __libc_start_main + 0xf3 (0x7f65fe02f083 in /lib/x86_64-linux-gnu/libc.so.6)
frame #11: <unknown function> + 0x58ee (0x56183fd188ee in ./build/src/interpolate_tensorrt)

That looked kind of similar to this issue and patching Torch-TensorRT with this PR makes the behavior exactly the same as in the second case (i.e. when using pytorch 1.13.1 and torch-tensorRT main).

Torch-TensorRT main (commit 861edd0)

• Torch-TensorRT Version (e.g. 1.0.0): main, commit 861edd0
• PyTorch Version (e.g. 1.0): 1.13.1
• CPU Architecture: x86_64
• OS (e.g., Linux): Linux, Ubuntu 20.04
• How you installed PyTorch (conda, pip, libtorch, source): pip for the python package used to generate the torchscript, source for the C++ dependency linked to Torch-TensorRT
• Build command you used (if compiling from source): Compiled with CMake
• Python version: 3.8.16
• CUDA version: 11.8
• GPU models and configuration: NVIDIA GeForce RTX 2070 SUPER, SM Capability: 7.5
• Any other relevant information:

When using torch-TensorRT main, the conversion just hangs for ever after

GRAPH: [Torch-TensorRT] - Torch-TensorRT.TorchScript Graph Lowering

Additional context

Interestingly, when using the tracing mechanism of pytorch to generate the torchscript, everything seems fine (I didn't check the results, but the conversion finishes properly).
Also, when scripting with pytorch 1.9, everything works fine 🤯

The thing I noticed is that pytorch changed slightly the torch.nn.interpolate API and I am wondering if that could explain (at least partially) the problem:

• Before pytorch 1.11: torch.nn.functional.interpolate(input, size=None, scale_factor=None, mode='nearest', align_corners=None, recompute_scale_factor=None)
• PyTorch 1.11 and more recent: torch.nn.functional.interpolate(input, size=None, scale_factor=None, mode='nearest', align_corners=None, recompute_scale_factor=None, antialias=False)

See the attached .zip file containing a python file to generate the torchscript.
upsample.zip

Let me know if you need more details to reproduce the problem. Thanks!

[gs-olive]

I've reproduced the error on main, and it is occurring on this line, where the operation torch::jit::EliminateExceptions does not complete. This is the source code for torch::jit::EliminateExceptions.

TensorRT/core/lowering/lowering.cpp

Line 106 in a245b86

torch::jit::EliminateExceptions(g);

The graph at that point is shown below, and has a prim::RaiseException type within a prim::If

Graph

graph(%X.1 : Tensor):
  %2 : bool = prim::Constant[value=0]()
  %3 : float[] = prim::Constant[value=[2., 2.]]()
  %4 : str = prim::Constant[value="bilinear"]() # /usr/local/lib/python3.8/dist-packages/torch/nn/modules/upsampling.py:156:66
  %5 : int = prim::Constant[value=5]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3956:22
  %6 : int = prim::Constant[value=3]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3952:22
  %7 : int = prim::Constant[value=4]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3949:76
  %8 : int = prim::Constant[value=2]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3880:24
  %9 : NoneType = prim::Constant() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3871:32
  %10 : str = prim::Constant[value="builtins.ValueError"]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3872:18
  %11 : str = prim::Constant[value="Input and scale_factor must have the same number of spatial dimensions, but got input with spatial dimensions of {} and scale_factor of shape {}. Please provide input tensor in (N, C, d1, d2, ...,dK) format and scale_factor in (s1, s2, ...,sK) format."]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3909:20
  %12 : str = prim::Constant[value="Got 3D input, but bilinear mode needs 4D input"]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3994:34
  %13 : str = prim::Constant[value="builtins.NotImplementedError"]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3994:14
  %14 : str = prim::Constant[value="Got 5D input, but bilinear mode needs 4D input"]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4004:34
  %15 : str = prim::Constant[value="Input Error: Only 3D, 4D and 5D input Tensors supported (got {}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | area | nearest-exact (got {})"]() # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4007:8
  %16 : int = prim::Constant[value=1]()
  %48 : int = prim::Constant[value=2]()
  %49 : str = prim::Constant[value="builtins.ValueError"]()
  %50 : float[] = prim::Constant[value=[2., 2.]]()
  %51 : str = prim::Constant[value="Input and scale_factor must have the same number of spatial dimensions, but got input with spatial dimensions of {} and scale_factor of shape {}. Please provide input tensor in (N, C, d1, d2, ...,dK) format and scale_factor in (s1, s2, ...,sK) format."]()
  %52 : int = prim::Constant[value=1]()
  %53 : NoneType = prim::Constant()
  %54 : int = prim::Constant[value=4]()
  %55 : str = prim::Constant[value="Input Error: Only 3D, 4D and 5D input Tensors supported (got {}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | area | nearest-exact (got {})"]()
  %56 : str = prim::Constant[value="bilinear"]()
  %57 : str = prim::Constant[value="Got 5D input, but bilinear mode needs 4D input"]()
  %58 : int = prim::Constant[value=5]()
  %59 : str = prim::Constant[value="builtins.NotImplementedError"]()
  %60 : str = prim::Constant[value="Got 3D input, but bilinear mode needs 4D input"]()
  %61 : int = prim::Constant[value=3]()
  %62 : bool = prim::Constant[value=0]()
  %63 : Tensor = prim::Uninitialized()
  %64 : int = aten::dim(%X.1) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3880:10
  %dim.2 : int = aten::sub(%64, %48) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3880:10
  %66 : bool = aten::ne(%48, %dim.2) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3907:15
   = prim::If(%66) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3907:12
    block0():
      %67 : int[] = aten::size(%X.1) # <string>:13:9
      %68 : int[] = aten::slice(%67, %48, %53, %52) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3909:26
      %69 : int[] = aten::list(%68) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3909:21
      %70 : str = aten::format(%51, %69, %50) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3909:20
       = prim::RaiseException(%70, %49) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3908:16
      -> ()
    block1():
      -> ()
  %71 : bool = aten::eq(%64, %54) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3979:7
  %out1.1 : Tensor = prim::If(%71) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3979:4
    block0():
      %73 : Tensor = aten::upsample_bilinear2d(%X.1, %53, %62, %50) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3983:15
      -> (%73)
    block1():
      %74 : bool = aten::eq(%64, %61) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3993:7
       = prim::If(%74) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3993:4
        block0():
           = prim::RaiseException(%60, %59) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3994:8
          -> ()
        block1():
          -> ()
      %75 : bool = aten::eq(%64, %58) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4003:7
       = prim::If(%75) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4003:4
        block0():
           = prim::RaiseException(%57, %59) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4004:8
          -> ()
        block1():
          -> ()
      %76 : str = aten::format(%55, %64, %56) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4007:8
       = prim::RaiseException(%76, %59) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4006:4
      -> (%63)
   = prim::If(%66) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3907:12
    block0():
      %77 : int[] = aten::size(%X.1) # <string>:13:9
      %78 : int[] = aten::slice(%77, %48, %53, %52) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3909:26
      %79 : int[] = aten::list(%78) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3909:21
      %80 : str = aten::format(%51, %79, %50) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3909:20
       = prim::RaiseException(%80, %49) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3908:16
      -> ()
    block1():
      -> ()
  %out2.1 : Tensor = prim::If(%71) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3979:4
    block0():
      %82 : Tensor = aten::upsample_bilinear2d(%X.1, %53, %62, %50) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3983:15
      -> (%82)
    block1():
      %83 : bool = aten::eq(%64, %61) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3993:7
       = prim::If(%83) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3993:4
        block0():
           = prim::RaiseException(%60, %59) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:3994:8
          -> ()
        block1():
          -> ()
      %84 : bool = aten::eq(%64, %58) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4003:7
       = prim::If(%84) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4003:4
        block0():
           = prim::RaiseException(%57, %59) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4004:8
          -> ()
        block1():
          -> ()
      %85 : str = aten::format(%55, %64, %56) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4007:8
       = prim::RaiseException(%85, %59) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:4006:4
      -> (%63)
  %out.1 : Tensor = aten::add(%out1.1, %out2.1, %16) # unrelated.py:24:14
  return (%out.1)

@bowang007 - this seems related to your work with exceptions and control flow, do you have any suggestions on this?

[gcuendet]

Thanks @gs-olive for having taken a look at this issue so quickly! Nice that you could reproduce it!
Following up on that, here is another simple Network (though slightly less trivial than the one above).

Network definition

import torch
import torch.nn as nn
import torch.nn.functional as F
import argparse


class Block(nn.Module):
    def __init__(self, in_channel, out_channel):
        super(Block, self).__init__()
        self.conv = nn.Conv2d(
            in_channel, out_channel, kernel_size=3, stride=1, padding=1, bias=False
        )
        self.norm = nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(
            kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False
        )

    def forward(self, x):
        out = self.conv(x)
        out = self.norm(out)
        out = self.relu(out)
        out = self.maxpool(out)
        return out


class Network(torch.nn.Module):
    def __init__(self, num_classes=2):
        super(Network, self).__init__()
        self.num_classes = num_classes

        self.block1 = Block(3, 32)
        self.block2 = Block(32, 64)

        self.upsample1 = nn.Upsample(
            scale_factor=2, mode="bilinear", align_corners=False
        )

        self.upsample2 = nn.Upsample(
            scale_factor=2, mode="bilinear", align_corners=False
        )

        self.conv = nn.Conv2d(64, num_classes, 1, bias=True)

    def forward(self, x):
        out = self.block1(x)
        out = self.block2(out)

        gclayer1 = self.upsample1(out)
        gclayer2 = self.upsample2(gclayer1)

        out = self.conv(gclayer2)
        return out

Interestingly, when converting a torchscript generated from that network using torch.jit.script on linux, the behaviour is the same as with the trivial network previously shared: using torch-tensorrt recent commit from main, it hangs (i.e. the operation torch::jit::EliminateExceptions does not complete).
Nevertheless, when converting the torchscript on windows using the same torch-tensorrt recent commit from main, it works! Note that it works both when the torchscript is generated on windows or on linux.

Both graphs, as printed by Torch-TensorRT when calling compile are included in archive.zip, as well as the mini_net.py network definition/torchscript generation file. In summary, here is a diff between the graphs, as printed by Torch-TensorRT:

Graph diff

--- graph_lin.py	2023-04-19 13:28:59.000000000 +0200
+++ graph_win.py	2023-04-19 13:28:52.000000000 +0200
@@ -1,16 +1,13 @@
 GRAPH: [Torch-TensorRT] - After freeze: graph(%self.1 : __torch__.___torch_mangle_0.Network,
       %x.1 : Tensor):
-  %224 : str[] = prim::Constant[value=["nearest", "area", "nearest-exact"]]()
   %115 : str = prim::Constant[value="bilinear"]() # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/modules/upsampling.py:153:66
   %114 : float = prim::Constant[value=2.]() # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/modules/upsampling.py:153:47
   %113 : int = prim::Constant[value=5]() # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3892:22
   %112 : int = prim::Constant[value=3]() # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3888:22
   %111 : int = prim::Constant[value=4]() # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3885:76
-  %109 : str = prim::Constant[value="align_corners option can only be set with the interpolating modes: linear | bilinear | bicubic | trilinear"]() # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3815:16
   %108 : str = prim::Constant[value="Got 3D input, but bilinear mode needs 4D input"]() # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3930:34
   %107 : str = prim::Constant[value="Got 5D input, but bilinear mode needs 4D input"]() # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3940:34
   %106 : str = prim::Constant[value="Input Error: Only 3D, 4D and 5D input Tensors supported (got {}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | area | nearest-exact (got {})"]() # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3943:8
-  %101 : str = prim::Constant[value="AssertionError: "]() # :0:0
   %223 : int[] = prim::Constant[value=[0, 0]]()
   %222 : int[] = prim::Constant[value=[2, 2]]()
   %221 : int[] = prim::Constant[value=[1, 1]]()
@@ -38,15 +35,6 @@
   %out1.1 : Tensor = aten::relu_(%out0.1) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:1440:17
   %out0.2 : Tensor = aten::max_pool2d(%out1.1, %222, %222, %223, %221, %self.block1.norm.training) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:797:11
   %117 : Tensor = prim::Uninitialized() # :0:0
-  %118 : bool? = prim::Uninitialized() # :0:0
-  %119 : bool = prim::Uninitialized() # :0:0
-  %121 : bool = aten::__contains__(%224, %115) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3812:7
-  %align_corners0.1 : bool? = prim::If(%121) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3812:4
-    block0():
-       = prim::RaiseException(%109, %self.block1.conv.bias) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3814:12
-      -> (%118)
-    block1():
-      -> (%self.block1.norm.training)
   %123 : int = aten::dim(%out0.2) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3822:10
   %dim.2 : int = aten::sub(%123, %18) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3822:10
   %scale_factors2.2 : float[] = prim::ListConstruct()
@@ -57,15 +45,7 @@
   %129 : bool = aten::eq(%123, %111) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3915:7
   %gclayer1.1 : Tensor = prim::If(%129) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3915:4
     block0():
-      %132 : bool = aten::__isnot__(%align_corners0.1, %self.block1.conv.bias) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3916:15
-      %align_corners6.1 : bool = prim::If(%132) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3916:8
-        block0():
-          %align_corners7.2 : bool = prim::unchecked_cast(%align_corners0.1)
-          -> (%align_corners7.2)
-        block1():
-           = prim::RaiseException(%101, %self.block1.conv.bias) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3916:8
-          -> (%119)
-      %135 : Tensor = aten::upsample_bilinear2d(%out0.2, %self.block1.conv.bias, %align_corners6.1, %scale_factors2.2) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3919:15
+      %135 : Tensor = aten::upsample_bilinear2d(%out0.2, %self.block1.conv.bias, %self.block1.norm.training, %scale_factors2.2) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3919:15
       -> (%135)
     block1():
       %137 : bool = aten::eq(%123, %112) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3929:7
@@ -85,12 +65,6 @@
       %143 : str = aten::format(%106, %123, %115) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3943:8
        = prim::RaiseException(%143, %self.block1.conv.bias) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3942:4
       -> (%117)
-  %align_corners0 : bool? = prim::If(%121) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3812:4
-    block0():
-       = prim::RaiseException(%109, %self.block1.conv.bias) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3814:12
-      -> (%118)
-    block1():
-      -> (%self.block1.norm.training)
   %167 : int = aten::dim(%gclayer1.1) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3822:10
   %dim.1 : int = aten::sub(%167, %18) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3822:10
   %scale_factors2.1 : float[] = prim::ListConstruct()
@@ -101,15 +75,7 @@
   %173 : bool = aten::eq(%167, %111) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3915:7
   %gclayer2.1 : Tensor = prim::If(%173) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3915:4
     block0():
-      %176 : bool = aten::__isnot__(%align_corners0, %self.block1.conv.bias) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3916:15
-      %align_corners6 : bool = prim::If(%176) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3916:8
-        block0():
-          %align_corners7.1 : bool = prim::unchecked_cast(%align_corners0)
-          -> (%align_corners7.1)
-        block1():
-           = prim::RaiseException(%101, %self.block1.conv.bias) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3916:8
-          -> (%119)
-      %179 : Tensor = aten::upsample_bilinear2d(%gclayer1.1, %self.block1.conv.bias, %align_corners6, %scale_factors2.1) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3919:15
+      %179 : Tensor = aten::upsample_bilinear2d(%gclayer1.1, %self.block1.conv.bias, %self.block1.norm.training, %scale_factors2.1) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3919:15
       -> (%179)
     block1():
       %181 : bool = aten::eq(%167, %112) # /home/gcuendet/.pyenv/versions/pytorch1.11/lib/python3.8/site-packages/torch/nn/functional.py:3929:7

Please let me know if I can provide more details to help solve this issue!

[gs-olive]

Thank you for the additional details - this is very helpful!

[gcuendet]

Hi @gs-olive
I have been working on the idea you described in #1842 . See this commit. Of course, this initial implementation is overly specific and only solves the case of upsample_bilinear2d but together with your proposed fix of torch::jit::EliminateExceptions and an additional small modification of torch-tensorRT custom exceptions elimination pass, is seems that I am able to successfully convert the two networks linked in this issue (the simple upsample, as well as the slightly less trivial one).

I would appreciate some feedback on this. Is my approach going in the right direction, w.r.t. what you had in mind when describing the #1842 issue?

[gcuendet]

Also, what do you (or @bowang007 maybe?) think of the changes I made to TensorRT/core/lowering/passes/exception_elimination.cpp ? The rationale is that instead of catching only something like:

 = prim::If(%5958)
  block0():
    = prim::RaiseException(%45)
   -> ()
  block1():
   -> ()

or

 = prim::If(%5958)
  block0():
   -> ()
  block1():
    = prim::RaiseException(%45)
   -> ()

you could also catch more complex blocks, given that:

1. the node is a prim::If, containing two blocks
2. both blocks do not return anything
3. the last node of at least one the block is a prim::RaiseException (instead of the first one currently)

What I've typically been observing in the scope of the Upsample investigation is something like:

 = prim::If(%5958)
  block0():
   -> ()
  block1():
    %191 : str = aten::format(%108, %171, %117)
   = prim::RaiseException(%191, %100)
   -> ()