From 2d8f03013483abe607bc52fe5598d077e2a96de4 Mon Sep 17 00:00:00 2001 From: Ilyabasharov Date: Sat, 26 Feb 2022 20:42:19 +0300 Subject: [PATCH 01/13] add avg_pool1d converter --- torch2trt/converters/avg_pool.py | 65 +++++++++++++++++++++++++++++--- 1 file changed, 60 insertions(+), 5 deletions(-) diff --git a/torch2trt/converters/avg_pool.py b/torch2trt/converters/avg_pool.py index 185af508..02eaea1a 100644 --- a/torch2trt/converters/avg_pool.py +++ b/torch2trt/converters/avg_pool.py @@ -2,6 +2,55 @@ from torch2trt.module_test import add_module_test +@tensorrt_converter('torch.nn.functional.avg_pool1d') +def convert_avg_pool1d(ctx): + # At the time of this implementation, TensorRT 8.x does not yet support avg pooling in 1D using `add_pooling_nd(...)`. + # As such, we use a workaround here, by unsqueezing another dimension into the input (thus transforming it from + # (N, C, L) to (N, C, L, 1)) so that we can use 2D max pooling across the last three dimensions. + + input = get_arg(ctx, 'input', pos=0, default=None) + input_trt = trt_(ctx.network, input) + output = ctx.method_return + + kernel_size = get_arg(ctx, 'kernel_size', pos=1, default=None) + stride = get_arg(ctx, 'stride', pos=2, default=None) + padding = get_arg(ctx, 'padding', pos=3, default=0) + ceil_mode = get_arg(ctx, 'ceil_mode', pos=4, default=False) + count_include_pad = get_arg(ctx, 'count_include_pad', pos=5, default=True) + + # Convert inputs to be 2d compatible as inputs will always be 1d. + kernel_size = (kernel_size, 1) + stride = kernel_size if not stride else (stride, 1) + padding = (padding, 0) + + # Shuffle layer to unsqueeze another dimension for 2D max pooling. + unsqueeze_layer = ctx.network.add_shuffle(input_trt) + set_layer_precision(ctx, unsqueeze_layer) + unsqueeze_layer.reshape_dims = tuple([*input_trt.shape, 1]) + unsqueeze_trt = unsqueeze_layer.get_output(0) + + # Use 2D max pooling here to fake 1D max pooling. + layer = ctx.network.add_pooling_nd( + input=unsqueeze_trt, + type=trt.PoolingType.AVERAGE, + window_size=kernel_size, + ) + set_layer_precision(ctx, layer) + layer.stride_nd = stride + layer.padding_nd = padding + layer.average_count_excludes_padding = not count_include_pad + + if ceil_mode: + layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP + + pooling_trt = layer.get_output(0) + + # Shuffle layer to squeeze out dimension that was just added for 2D max pooling so return is still in 1D. + squeeze_layer = ctx.network.add_shuffle(pooling_trt) + set_layer_precision(ctx, squeeze_layer) + squeeze_layer.reshape_dims = tuple(pooling_trt.shape[:-1]) + output._trt = squeeze_layer.get_output(0) + @tensorrt_converter("torch.nn.functional.avg_pool2d", enabled=trt_version() < '7.0') def convert_avg_pool2d(ctx): # parse args @@ -83,12 +132,14 @@ def convert_avg_pool_trt7(ctx): layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP output._trt = layer.get_output(0) - - + + @add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 4, 6)]) @add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 5, 7)]) def test_avg_pool2d_without_ceil_mode(): - return torch.nn.AvgPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=False) + return torch.nn.AvgPool2d( + kernel_size=3, stride=2, padding=1, ceil_mode=False + ) @add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 4, 6)]) @@ -102,10 +153,14 @@ def test_avg_pool2d_with_ceil_mode(): @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 4, 4, 6)], enabled=trt_version() >= '7.0') @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 3, 5, 7)], enabled=trt_version() >= '7.0') def test_avg_pool3d_without_ceil_mode_trt7(): - return torch.nn.AvgPool3d(kernel_size=3, stride=2, padding=1, ceil_mode=False) + return torch.nn.AvgPool3d( + kernel_size=3, stride=2, padding=1, ceil_mode=False + ) @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 4, 4, 6)], enabled=trt_version() >= '7.0') @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 3, 5, 7)], enabled=trt_version() >= '7.0') def test_avg_pool3d_with_ceil_mode_trt7(): - return torch.nn.AvgPool3d(kernel_size=3, stride=2, padding=1, ceil_mode=True, count_include_pad=False) # TRT does not support ceil_mode=True && count_include_pad=True + return torch.nn.AvgPool3d( + kernel_size=3, stride=2, padding=1, ceil_mode=True, count_include_pad=False + ) # TRT does not support ceil_mode=True && count_include_pad=True From 34f1c009517ce0d6a0e251dbfdfa68a3b9c7e9ad Mon Sep 17 00:00:00 2001 From: Ilyabasharov Date: Sat, 26 Feb 2022 20:43:16 +0300 Subject: [PATCH 02/13] Add expand_as converter --- torch2trt/converters/expand.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch2trt/converters/expand.py b/torch2trt/converters/expand.py index e0d07540..2a22a55d 100644 --- a/torch2trt/converters/expand.py +++ b/torch2trt/converters/expand.py @@ -2,10 +2,10 @@ from torch2trt.module_test import add_module_test +@tensorrt_converter('torch.Tensor.expand_as') @tensorrt_converter('torch.Tensor.expand') def convert_expand(ctx): input = ctx.method_args[0] - sizes = ctx.method_args[1:] output = ctx.method_return inshape = tuple(input.shape)[1:] # exclude batch From 45bff6c0ae950a7a98f7d825c45eb5047bd970ad Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Sun, 27 Feb 2022 02:54:16 +0300 Subject: [PATCH 03/13] Update module_test.py --- torch2trt/module_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/torch2trt/module_test.py b/torch2trt/module_test.py index fb158fe8..3a316407 100644 --- a/torch2trt/module_test.py +++ b/torch2trt/module_test.py @@ -1,7 +1,3 @@ -import torch -import torchvision - - class ModuleTest(object): def __init__(self, module_fn, dtype, device, input_shapes, **torch2trt_kwargs): self.module_fn = module_fn From 688546075ef7d3c3a17211dad23b27d1fa603b43 Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Sat, 5 Mar 2022 22:01:28 +0400 Subject: [PATCH 04/13] Update setup.py --- setup.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/setup.py b/setup.py index 3df5cc97..86c671dd 100644 --- a/setup.py +++ b/setup.py @@ -5,6 +5,10 @@ from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension from packaging import version +REQUIREMENTS_PATH = 'requirements.txt' + +with open(REQUIREMENTS_PATH, 'r') as file: + required_libraries = file.read().splitlines() def trt_inc_dir(): return "/usr/include/aarch64-linux-gnu" @@ -55,5 +59,6 @@ def trt_lib_dir(): packages=find_packages(exclude=exclude_dir), ext_package='torch2trt', ext_modules=ext_modules, + install_requires=required_libraries, cmdclass={'build_ext': BuildExtension} ) From e257927042c6545e8e0f8950123c9771f5b04005 Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Sat, 5 Mar 2022 22:01:57 +0400 Subject: [PATCH 05/13] Create requirements.txt --- requirements.txt | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 requirements.txt diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 00000000..2100b161 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +packaging +termcolor From 7eaf1cff45bbe9e117d66fb4195615cd4b935080 Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Sat, 5 Mar 2022 22:03:48 +0400 Subject: [PATCH 06/13] python3 fix --- scripts/build_contrib.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/build_contrib.sh b/scripts/build_contrib.sh index 62462176..8d4de109 100755 --- a/scripts/build_contrib.sh +++ b/scripts/build_contrib.sh @@ -13,7 +13,7 @@ pushd /tmp/TensorRT git sparse-checkout set /tools/pytorch-quantization/ git apply --reject --whitespace=fix pytorch_nvidia_quantization.patch cd tools/pytorch-quantization/ - python setup.py install + python3 setup.py install popd pushd $parentdir From 31738c29873e0cf98f24382a5567fef8beb1adca Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Sat, 5 Mar 2022 22:44:02 +0400 Subject: [PATCH 07/13] Update requirements.txt --- requirements.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 2100b161..f08cca14 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1 @@ -packaging termcolor From c00a45f62c99b9c1d9173dd3e8ee6e54858dc261 Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Wed, 9 Mar 2022 15:52:25 +0400 Subject: [PATCH 08/13] Update __init__.py --- torch2trt/converters/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch2trt/converters/__init__.py b/torch2trt/converters/__init__.py index 5ffe38b4..bd36e15e 100644 --- a/torch2trt/converters/__init__.py +++ b/torch2trt/converters/__init__.py @@ -32,7 +32,7 @@ from .expand import * from .floordiv import * from .gelu import * -from .getitem import * +#from .getitem import * from .group_norm import * from .identity import * from .instance_norm import * From 6cab8a98a0cddaf673d8b5368dab992262d0b875 Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Thu, 10 Mar 2022 15:44:13 +0400 Subject: [PATCH 09/13] Update __init__.py --- torch2trt/converters/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch2trt/converters/__init__.py b/torch2trt/converters/__init__.py index bd36e15e..5ffe38b4 100644 --- a/torch2trt/converters/__init__.py +++ b/torch2trt/converters/__init__.py @@ -32,7 +32,7 @@ from .expand import * from .floordiv import * from .gelu import * -#from .getitem import * +from .getitem import * from .group_norm import * from .identity import * from .instance_norm import * From 47b77cf0e4403210c09a6d855ab990febe71ed7e Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Wed, 16 Mar 2022 02:02:22 +0300 Subject: [PATCH 10/13] Update build_contrib.sh --- scripts/build_contrib.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/build_contrib.sh b/scripts/build_contrib.sh index 8d4de109..29f76541 100755 --- a/scripts/build_contrib.sh +++ b/scripts/build_contrib.sh @@ -13,11 +13,11 @@ pushd /tmp/TensorRT git sparse-checkout set /tools/pytorch-quantization/ git apply --reject --whitespace=fix pytorch_nvidia_quantization.patch cd tools/pytorch-quantization/ - python3 setup.py install + sudo python3 setup.py install popd pushd $parentdir - python3 setup.py install --plugins --contrib + sudo python3 setup.py install --plugins --contrib popd From 4b00a99564a912b8456ef9dbf724daf975c8b7a1 Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Thu, 31 Mar 2022 23:00:01 +0300 Subject: [PATCH 11/13] Create zeros.py --- torch2trt/converters/zeros.py | 66 +++++++++++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 torch2trt/converters/zeros.py diff --git a/torch2trt/converters/zeros.py b/torch2trt/converters/zeros.py new file mode 100644 index 00000000..b8209576 --- /dev/null +++ b/torch2trt/converters/zeros.py @@ -0,0 +1,66 @@ +from torch2trt.torch2trt import * +from torch2trt.module_test import add_module_test + + +def _set_layer_precision(ctx, layer): + # Supported TRT precisions as given by torch2trt_kwargs. + INT8_MODE = "int8_mode" + FP16_MODE = "fp16_mode" + + # Check that args exist as expected in torch2trt_kwargs. + trt_kwargs = ctx.torch2trt_kwargs + assert INT8_MODE in trt_kwargs + assert FP16_MODE in trt_kwargs + + is_int8 = trt_kwargs.get(INT8_MODE, False) + is_fp16 = trt_kwargs.get(FP16_MODE, False) + + if is_int8: + layer.precision = trt.int8 + layer.set_output_type(0, trt.int8) + elif is_fp16: + layer.precision = trt.float16 + layer.set_output_type(0, trt.float16) + + +@tensorrt_converter('torch.zeros') +def convert_zeros(ctx): + tensor = ctx.method_return + + # Implementation copied from add_trt_constant. + shape = tuple(tensor.shape[1:]) + array = tensor[0].detach().cpu().numpy() + layer = ctx.network.add_constant(shape, array) + + _set_layer_precision(ctx, layer) + + tensor._trt = layer.get_output(0) + + +class Zeros(torch.nn.Module): + def __init__(self, *size): + super().__init__() + self.size = size + + def forward(self, x): + return x + torch.zeros(*self.size, device=torch.device('cuda')) + + +@add_module_test(torch.float32, torch.device('cuda'), [(1, 2, 3, 4)]) +def test_zeros(): + return Zeros((1, 2, 3, 4)) + + +@add_module_test(torch.float32, torch.device('cuda'), [(1, 2, 3, 4)]) +def test_zeros_var_args(): + return Zeros(1, 2, 3, 4) + + +@add_module_test(torch.float32, torch.device('cuda'), [(1, 2, 3, 4)], fp16_mode=True) +def test_zeros_fp16_mode(): + return Zeros(1, 2, 3, 4) + + +@add_module_test(torch.float32, torch.device('cuda'), [(1, 2, 3, 4)], int8_mode=True) +def test_zeros_int8_mode(): + return Zeros(1, 2, 3, 4) From 02fd76ab21a7237642c3c5bb9695e979f3c55d8f Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Thu, 31 Mar 2022 23:10:43 +0300 Subject: [PATCH 12/13] Update __init__.py --- torch2trt/converters/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torch2trt/converters/__init__.py b/torch2trt/converters/__init__.py index 5ffe38b4..7413f3c7 100644 --- a/torch2trt/converters/__init__.py +++ b/torch2trt/converters/__init__.py @@ -69,3 +69,4 @@ from .transpose import * from .unary import * from .view import * +form .zeros import * From 76bf298b3da408509665e23e2494922b131afb10 Mon Sep 17 00:00:00 2001 From: Ilya Basharov Date: Thu, 31 Mar 2022 23:17:30 +0300 Subject: [PATCH 13/13] Update __init__.py --- torch2trt/converters/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch2trt/converters/__init__.py b/torch2trt/converters/__init__.py index 7413f3c7..a9fa49df 100644 --- a/torch2trt/converters/__init__.py +++ b/torch2trt/converters/__init__.py @@ -69,4 +69,4 @@ from .transpose import * from .unary import * from .view import * -form .zeros import * +from .zeros import *