Support trt 8 matmul + merge from upstream 18/10

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 ## [Master]
 
+- Added support for TensorRT 8
+
+## [0.3.0] - 07/15/2021
+
+- Added converter for ``torch.nn.functional.adaptive_avg_pool3d``
+- Added converter for ``torch.nn.functional.adaptive_max_pool3d``
+- Added converter for ``torch.maxpool3d`` and ``torch.nn.functional.max_pool3d``
 - Added Quantization Aware Training (QAT) workflow to contrib
 - Added converter for ``torch.roll``
 - Added converter for ``torch.nn.functional.layer_norm``

docker/21-06/Dockerfile

@@ -0,0 +1,4 @@
+FROM nvcr.io/nvidia/pytorch:21.06-py3
+
+
+RUN pip3 install termcolor

docker/21-06/build.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+docker build -t torch2trt:21-06 -f $(pwd)/docker/21-06/Dockerfile .

docker/21-06/run.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+
+docker run --gpus all -it --rm -v $(pwd):/torch2trt torch2trt:21-06 

docker/21-09/Dockerfile

@@ -0,0 +1,4 @@
+FROM nvcr.io/nvidia/pytorch:21.09-py3
+
+
+RUN pip3 install termcolor

docker/21-09/build.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+docker build -t torch2trt:21-09 -f $(pwd)/docker/21-09/Dockerfile .

docker/21-09/run.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+
+docker run --gpus all -it --rm -v $(pwd):/torch2trt torch2trt:21-09

examples/contrib/quantization_aware_training/README.md

@@ -25,27 +25,9 @@ RUN add-apt-repository ppa:git-core/ppa && \
 
 RUN pip install termcolor graphviz
 
-## If you have followed instructions on main README.md file to install torch2trt using scripts/build_contrib.sh
-## You dont require rest of the steps
-
-RUN git clone https://github.com/NVIDIA/TensorRT.git /sw/TensorRT/
-
-##Make sure that patch file is under the same folder where dockerfile is being called
-
-ADD pytorch_nvidia_quantization.patch /sw/TensorRT
-
-RUN cd /sw/TensorRT/ && \
-    git sparse-checkout init --cone && \
-    git sparse-checkout set /tools/pytorch-quantization/ && \
-    git apply --reject --whitespace=fix pytorch_nvidia_quantization.patch && \
-    cd tools/pytorch-quantization/ && \
-    python setup.py install 
-
-RUN git clone https://github.com/NVIDIA-AI-IOT/torch2trt.git /sw/TensorRT/ && \
-    cd /sw/TensorRT/ && \
-    git fetch origin pull/514/head:PR514 && \
-    git checkout PR514 && \
-    python setup.py install --plugins
+RUN git clone https://github.com/NVIDIA-AI-IOT/torch2trt.git /sw/torch2trt/ && \
+    cd /sw/torch2trt/scripts && \
+	bash build_contrib.sh
 
 ```
 

setup.py

@@ -41,7 +41,7 @@ def trt_lib_dir():
 
 setup(
     name='torch2trt',
-    version='0.2.0',
+    version='0.3.0',
     description='An easy to use PyTorch to TensorRT converter',
     packages=find_packages(exclude=exclude_dir),
     ext_package='torch2trt',

torch2trt/converters/AdaptiveAvgPool3d.py

@@ -0,0 +1,48 @@
+from torch2trt.torch2trt import *
+from torch2trt.module_test import add_module_test
+
+
+@tensorrt_converter(
+    "torch.nn.AdaptiveAvgPool3d.forward", enabled=trt_version() >= "7.0"
+)
+def convert_AdaptiveAvgPool3d(ctx):
+    module = ctx.method_args[0]
+    input = ctx.method_args[1]
+    output = ctx.method_return
+
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
+
+    output_size = module.output_size
+    if not isinstance(output_size, tuple):
+        output_size = (output_size,) * 3
+
+    stride = (
+        input_trt.shape[-3] // output_size[-3],
+        input_trt.shape[-2] // output_size[-2],
+        input_trt.shape[-1] // output_size[-1],
+    )
+
+    kernel_size = stride
+    layer = ctx.network.add_pooling_nd(
+        input=input_trt,
+        type=trt.PoolingType.AVERAGE,
+        window_size=kernel_size,
+    )
+    layer.stride_nd = stride
+
+    output._trt = layer.get_output(0)
+
+
+@add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 16, 224, 224)])
+def test_AdaptiveAvgPool3d_1x1x1():
+    return torch.nn.AdaptiveAvgPool3d((1, 1, 1))
+
+
+@add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 16, 224, 224)])
+def test_AdaptiveAvgPool3d_2x2x2():
+    return torch.nn.AdaptiveAvgPool3d((2, 2, 2))
+
+
+@add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 16, 224, 224)])
+def test_AdaptiveAvgPool3d_3x3x3():
+    return torch.nn.AdaptiveAvgPool3d((3, 3, 3))

torch2trt/converters/BatchNorm3d.py

@@ -0,0 +1,23 @@
+from torch2trt.torch2trt import *
+from torch2trt.module_test import add_module_test
+
+
+@tensorrt_converter("torch.nn.BatchNorm3d.forward", enabled=trt_version() < "7.0")
+def convert_BatchNorm3d(ctx):
+    module = ctx.method_args[0]
+    input = ctx.method_args[1]
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
+    output = ctx.method_return
+
+    scale = module.weight.detach().cpu().numpy() / np.sqrt(
+        module.running_var.detach().cpu().numpy() + module.eps
+    )
+    bias = (
+        module.bias.detach().cpu().numpy()
+        - module.running_mean.detach().cpu().numpy() * scale
+    )
+    power = np.ones_like(scale)
+
+    layer = ctx.network.add_scale(input_trt, trt.ScaleMode.CHANNEL, bias, scale, power)
+
+    output._trt = layer.get_output(0)

torch2trt/converters/__init__.py

@@ -7,6 +7,7 @@
 from .AdaptiveAvgPool2d import *
 from .BatchNorm1d import *
 from .BatchNorm2d import *
+from .clone import *
 from .conv_functional import *
 from .Conv import *
 from .Conv1d import *
@@ -17,7 +18,9 @@
 from .LogSoftmax import *
 from .activation import *
 from .adaptive_avg_pool2d import *
+from .adaptive_avg_pool3d import *
 from .adaptive_max_pool2d import *
+from .adaptive_max_pool3d import *
 from .add import *
 from .avg_pool import *
 from .batch_norm import *
@@ -35,8 +38,11 @@
 from .instance_norm import *
 from .interpolate import *
 from .layer_norm import *
+from .matmul import *
 from .max import *
+from .max_pool1d import *
 from .max_pool2d import *
+from .max_pool3d import *
 from .mean import *
 from .min import *
 from .mod import *

torch2trt/converters/adaptive_avg_pool3d.py

@@ -0,0 +1,11 @@
+from torch2trt.torch2trt import *
+from .AdaptiveAvgPool3d import *
+
+
+@tensorrt_converter("torch.nn.functional.adaptive_avg_pool3d")
+def convert_adaptive_avg_pool3d(ctx):
+    ctx.method_args = (
+        torch.nn.AdaptiveAvgPool3d(ctx.method_args[1]),
+        ctx.method_args[0],
+    )
+    convert_AdaptiveAvgPool3d(ctx)

torch2trt/converters/adaptive_max_pool3d.py

@@ -0,0 +1,41 @@
+from torch2trt.torch2trt import *
+from torch2trt.module_test import add_module_test
+
+
+@tensorrt_converter("torch.nn.functional.adaptive_max_pool3d")
+def convert_adaptive_max_pool3d(ctx):
+    input = ctx.method_args[0]
+    output = ctx.method_return
+
+    output_size = ctx.method_args[1]
+    if isinstance(output_size, int):
+        output_size = (output_size,) * 3
+
+    stride = (
+        input._trt.shape[-3] // output_size[-3],
+        input._trt.shape[-2] // output_size[-2],
+        input._trt.shape[-1] // output_size[-1],
+    )
+
+    kernel_size = stride
+    layer = ctx.network.add_pooling_nd(
+        input=input._trt, type=trt.PoolingType.MAX, window_size=kernel_size
+    )
+    layer.stride_nd = stride
+
+    output._trt = layer.get_output(0)
+
+
+@add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 16, 224, 224)])
+def test_adaptive_max_pool3d_1x1x1():
+    return torch.nn.AdaptiveMaxPool3d((1, 1, 1))
+
+
+@add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 16, 224, 224)])
+def test_adaptive_max_pool3d_2x2x2():
+    return torch.nn.AdaptiveMaxPool3d((2, 2, 2))
+
+
+@add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 16, 224, 224)])
+def test_adaptive_max_pool3d_3x3x3():
+    return torch.nn.AdaptiveMaxPool3d((3, 3, 3))

torch2trt/converters/clone.py

@@ -0,0 +1,83 @@
+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.clone')
+@tensorrt_converter('torch.Tensor.clone')
+def convert_clone(ctx):
+    input = ctx.method_args[0]
+    input_trt = trt_(ctx.network, input)
+
+    # Clone by making identity layer.
+    layer = ctx.network.add_identity(input_trt)
+    _set_layer_precision(ctx, layer)
+
+    output = ctx.method_return
+    output._trt = layer.get_output(0)
+
+
+class Clone(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.clone()
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 64, 64)])
+def test_clone_basic():
+    return Clone()
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 64, 64)], fp16_mode=True)
+def test_clone_fp16_mode():
+    return Clone()
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 64, 64)], int8_mode=True)
+def test_clone_int8_mode():
+    return Clone()
+
+
+class TorchClone(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.clone(x)
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 64, 64)])
+def test_torch_clone_basic():
+    return TorchClone()
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 64, 64)], fp16_mode=True)
+def test_torch_clone_fp16_mode():
+    return TorchClone()
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 64, 64)], int8_mode=True)
+def test_torch_clone_int8_mode():
+    return TorchClone()

torch2trt/converters/dummy_converters.py

@@ -3,7 +3,7 @@
 
 def is_private(method):
     method = method.split('.')[-1]  # remove prefix
-    return method[0] == '_' and method[1] is not '_'
+    return method[0] == '_' and method[1] != '_'
 
 def is_function_type(method):
     fntype =  eval(method + '.__class__.__name__')
@@ -34,4 +34,4 @@ def warn_method(ctx):
 @tensorrt_converter('torch.Tensor.dim', is_real=False)
 @tensorrt_converter('torch.Tensor.size', is_real=False)
 def dont_warn(ctx):
-    pass
+    pass