add allocator

.isort.cfg

@@ -1,2 +1,2 @@
 [settings]
-known_third_party = distutils,graphviz,numpy,packaging,setuptools,tensorrt,termcolor,torch,torchvision
+known_third_party = distutils,graphviz,mmdeploy,numpy,packaging,setuptools,tensorrt,torch,torchvision

tests/test_convert.py

@@ -0,0 +1,26 @@
+import torch
+from torch2trt_dynamic import TRTModule, module2trt
+from torchvision.models import resnet18
+
+
+def test_convert(tmp_path):
+    model = resnet18().cuda().eval()
+
+    trt_model = module2trt(
+        model,
+        args=[torch.rand(1, 3, 224, 224).cuda()],
+    )
+
+    model_path = tmp_path / 'tmp.pth'
+    torch.save(trt_model.state_dict(), model_path)
+    assert model_path.exists()
+
+    trt_model = TRTModule()
+    trt_model.load_state_dict(torch.load(model_path))
+
+    x = torch.rand(1, 3, 224, 224).cuda()
+    with torch.no_grad():
+        y = model(x)
+        y_trt = trt_model(x)
+
+    torch.testing.assert_close(y, y_trt)

torch2trt_dynamic/converters/Conv2d.py

@@ -34,15 +34,15 @@ def convert_Conv2d(ctx):
     if module.bias is not None:
         bias = module.bias.detach().cpu().numpy()
 
-    layer = ctx.network.add_convolution(
+    layer = ctx.network.add_convolution_nd(
         input=input_trt,
         num_output_maps=module.out_channels,
         kernel_shape=kernel_size,
         kernel=kernel,
         bias=bias)
-    layer.stride = stride
-    layer.padding = padding
-    layer.dilation = dilation
+    layer.stride_nd = stride
+    layer.padding_nd = padding
+    layer.dilation_nd = dilation
 
     if module.groups is not None:
         layer.num_groups = module.groups

torch2trt_dynamic/converters/Linear.py

@@ -21,7 +21,7 @@ def convert_Linear(ctx):
     if module.bias is not None:
         bias = module.bias.detach().cpu().numpy()
 
-    layer = ctx.network.add_convolution(
+    layer = ctx.network.add_convolution_nd(
         input=layer.get_output(0),
         num_output_maps=module.out_features,
         kernel_shape=(1, 1),

torch2trt_dynamic/converters/__init__.py

@@ -7,10 +7,9 @@
 
 # supported converters will override dummy converters
 
+from . import add  # noqa: F401
 from .activation import (convert_elu, convert_leaky_relu, convert_selu,
                          convert_softplus, convert_softsign)
-from .add import (convert_add, test_add_basic, test_add_iadd,
-                  test_add_radd_float, test_add_radd_int, test_add_torchadd)
 from .addcmul import convert_addcmul, test_addcmul
 from .arange import convert_arange
 from .argmax import convert_argmax
@@ -123,11 +122,6 @@
     'convert_softsign',
     'convert_softplus',
 ]
-# add
-__all__ += [
-    'convert_add', 'test_add_basic', 'test_add_iadd', 'test_add_radd_float',
-    'test_add_radd_int', 'test_add_torchadd'
-]
 # addcmul
 __all__ += ['convert_addcmul', 'test_addcmul']
 # arange

torch2trt_dynamic/converters/max_pool2d.py

@@ -32,11 +32,11 @@ def convert_max_pool2d(ctx):
     if not isinstance(padding, tuple):
         padding = (padding, ) * 2
 
-    layer = ctx.network.add_pooling(
+    layer = ctx.network.add_pooling_nd(
         input=input_trt, type=trt.PoolingType.MAX, window_size=kernel_size)
 
-    layer.stride = stride
-    layer.padding = padding
+    layer.stride_nd = stride
+    layer.padding_nd = padding
 
     if ceil_mode:
         layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP