delete the old file

Author: cehongwang

tests/py/dynamo/models/test_export_kwargs_serde.py

@@ -10,9 +10,8 @@
 import torch_tensorrt as torchtrt
 import torchvision.models as models
 from torch import nn
+from torch_tensorrt.dynamo._compiler import convert_module_to_trt_engine
 from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity
-from transformers import BertModel
-from transformers.utils.fx import symbolic_trace as transformers_trace
 
 assertions = unittest.TestCase()
 
@@ -60,8 +59,6 @@ def forward(self, x, b=5, c=None, d=None):
         "min_block_size": 1,
         "ir": "dynamo",
     }
-    # TODO: Support torchtrt.compile
-    # trt_mod = torchtrt.compile(model, **compile_spec)
 
     exp_program = torch.export.export(model, args=tuple(args), kwargs=kwargs)
     trt_gm = torchtrt.dynamo.compile(exp_program, **compile_spec)
@@ -76,3 +73,50 @@ def forward(self, x, b=5, c=None, d=None):
     torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwargs_inputs=kwargs)
     # Clean up model env
     torch._dynamo.reset()
+
+
+def test_custom_model_compile_engine():
+    class net(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv1 = nn.Conv2d(3, 12, 3, padding=1)
+            self.bn = nn.BatchNorm2d(12)
+            self.conv2 = nn.Conv2d(12, 12, 3, padding=1)
+            self.fc1 = nn.Linear(12 * 56 * 56, 10)
+
+        def forward(self, x, b=5, c=None, d=None):
+            x = self.conv1(x)
+            x = F.relu(x)
+            x = self.bn(x)
+            x = F.max_pool2d(x, (2, 2))
+            x = self.conv2(x)
+            x = F.relu(x)
+            x = F.max_pool2d(x, (2, 2))
+            x = torch.flatten(x, 1)
+            x = x + b
+            if c is not None:
+                x = x * c
+            if d is not None:
+                x = x - d["value"]
+            return self.fc1(x)
+
+    model = net().eval().to("cuda")
+    args = [torch.rand((1, 3, 224, 224)).to("cuda")]
+    kwargs = {
+        "b": torch.tensor(6).to("cuda"),
+        "d": {"value": torch.tensor(8).to("cuda")},
+    }
+
+    compile_spec = {
+        "inputs": args,
+        "kwarg_inputs": kwargs,
+        "device": torchtrt.Device("cuda:0"),
+        "enabled_precisions": {torch.float},
+        "pass_through_build_failures": True,
+        "optimization_level": 1,
+        "min_block_size": 1,
+        "ir": "dynamo",
+    }
+
+    exp_program = torch.export.export(model, args=tuple(args), kwargs=kwargs)
+    engine = convert_module_to_trt_engine(exp_program, **compile_spec)