Merge branch 'hyperfraise-3d'

Author: John Welsh

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## [Master]
 
+- 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``

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

@@ -17,7 +17,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 *
@@ -37,6 +39,7 @@
 from .layer_norm import *
 from .max 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/max_pool3d.py

@@ -0,0 +1,55 @@
+from torch2trt.torch2trt import *
+from torch2trt.module_test import add_module_test
+
+
+@tensorrt_converter("torch.nn.functional.max_pool3d")
+@tensorrt_converter("torch.max_pool3d")
+def convert_max_pool3d(ctx):
+    # parse args
+    input = get_arg(ctx, "input", pos=0, default=None)
+    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)
+    dilation = get_arg(ctx, "dilation", pos=4, default=1)
+    ceil_mode = get_arg(ctx, "ceil_mode", pos=5, default=False)
+
+    # get input trt tensor (or create constant if it doesn't exist)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
+
+    output = ctx.method_return
+
+    # get kernel size
+    if not isinstance(kernel_size, tuple):
+        kernel_size = (kernel_size,) * 3
+
+    # get stride
+    if not isinstance(stride, tuple):
+        stride = (stride,) * 3
+
+    # get padding
+    if not isinstance(padding, tuple):
+        padding = (padding,) * 3
+
+    layer = ctx.network.add_pooling_nd(
+        input=input_trt, type=trt.PoolingType.MAX, window_size=kernel_size
+    )
+
+    layer.stride_nd = stride
+    layer.padding_nd = padding
+
+    if ceil_mode:
+        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, 7)])
+@add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 5, 7, 8)])
+def test_MaxPool3d_without_ceil_mode():
+    return torch.nn.MaxPool3d(kernel_size=3, stride=2, padding=1, ceil_mode=False)
+
+
+@add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 4, 6, 7)])
+@add_module_test(torch.float32, torch.device("cuda"), [(1, 3, 5, 7, 8)])
+def test_MaxPool3d_with_ceil_mode():
+    return torch.nn.MaxPool3d(kernel_size=3, stride=2, padding=1, ceil_mode=True)