test: Add end-to-end tests for select models in FX2TRT

- Designed to reflect `test_models.py` which evaluates the TorchScript
path on a few key models, `test_compile_e2e.py` tests the
ease-of-compilation using fx2trt by porting over the test cases and
using "ir=fx"
- Evaluate functionality using cosine similarity metric
- Small fix in failing test within dispatch tracer
- Enable all tracer tests in noxfile
- BERT/Transformer tests to come

Author: gs-olive

noxfile.py

@@ -269,12 +269,12 @@ def run_fx_quant_tests(session):
 def run_fx_tracer_tests(session):
     print("Running FX Tracer tests")
     session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test"))
-    # skipping a test since it depends on torchdynamo
-    # Enable this test once NGC moves to latest pytorch which has dynamo integrated.
     tests = [
         "tracer/test_acc_shape_prop.py",
         "tracer/test_acc_tracer.py",
-        # "tracer/test_dispatch_tracer.py"
+        "tracer/test_dispatch_tracer.py"
+        "tracer/test_compile_e2e.py"
+        "tracer/test_resnet.py",
     ]
     for test in tests:
         if USE_HOST_DEPS:

py/torch_tensorrt/fx/test/tracer/test_compile_e2e.py

@@ -0,0 +1,106 @@
+import unittest
+import torch_tensorrt as torchtrt
+import torch
+import torchvision.models as models
+import timm
+
+
+COS_SIM_THRESHOLD = 0.99
+
+
+def cosine_similarity_custom(trt_out, torch_out):
+    torch.nn.functional.cosine_similarity(
+        trt_out.flatten(), torch_out.flatten(), dim=0, eps=1e-6
+    )
+
+
+class TestCompileE2E(unittest.TestCase):
+    def test_resnet18_fx(self):
+        self.model = models.resnet18(pretrained=True).eval().to("cuda")
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda")
+
+        compile_spec = {
+            "inputs": [self.input],
+            "enabled_precisions": {torch.float},
+        }
+
+        trt_mod = torchtrt.compile(self.model, ir="fx", **compile_spec)
+        cos_sim = cosine_similarity_custom(
+            self.model(self.input),
+            trt_mod(self.input),
+            dim=0,
+            eps=1e-6,
+        )
+        self.assertTrue(
+            cos_sim > COS_SIM_THRESHOLD,
+            msg=f"Resnet50 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COS_SIM_THRESHOLD}",
+        )
+
+    def test_mobilenet_v2_fx(self):
+        self.model = models.mobilenet_v2(pretrained=True).eval().to("cuda")
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda")
+
+        compile_spec = {
+            "inputs": [self.input],
+            "enabled_precisions": {torch.float},
+        }
+
+        trt_mod = torchtrt.compile(self.model, ir="fx", **compile_spec)
+        cos_sim = cosine_similarity_custom(
+            self.model(self.input),
+            trt_mod(self.input),
+            dim=0,
+            eps=1e-6,
+        )
+        self.assertTrue(
+            cos_sim > COS_SIM_THRESHOLD,
+            msg=f"Mobilenet v2 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COS_SIM_THRESHOLD}",
+        )
+
+    def test_efficientnet_b0_fx(self):
+        self.model = (
+            timm.create_model("efficientnet_b0", pretrained=True).eval().to("cuda")
+        )
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda")
+
+        compile_spec = {
+            "inputs": [self.input],
+            "enabled_precisions": {torch.float},
+        }
+
+        trt_mod = torchtrt.compile(self.model, ir="fx", **compile_spec)
+        cos_sim = cosine_similarity_custom(
+            self.model(self.input),
+            trt_mod(self.input),
+            dim=0,
+            eps=1e-6,
+        )
+        self.assertTrue(
+            cos_sim > COS_SIM_THRESHOLD,
+            msg=f"EfficientNet-B0 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COS_SIM_THRESHOLD}",
+        )
+
+    def test_resnet18_half_fx(self):
+        self.model = models.resnet18(pretrained=True).eval().to("cuda").half()
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda").half()
+
+        compile_spec = {
+            "inputs": [self.input],
+            "enabled_precisions": {torch.half},
+        }
+
+        trt_mod = torchtrt.compile(self.model, ir="fx", **compile_spec)
+        cos_sim = cosine_similarity_custom(
+            self.model(self.input),
+            trt_mod(self.input),
+            dim=0,
+            eps=1e-6,
+        )
+        self.assertTrue(
+            cos_sim > COS_SIM_THRESHOLD,
+            msg=f"Resnet18 Half TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COS_SIM_THRESHOLD}",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

py/torch_tensorrt/fx/test/tracer/test_dispatch_tracer.py

@@ -237,6 +237,7 @@ def fake_signature(fn, nargs):
             gm = make_fx(functionalize(fake_signature(gm, nargs)))(*example_inputs)
             return gm
 
+        torchdynamo.reset()
         optimized_mod = torchdynamo.optimize(
             compile_dispatch,
             nopython=True,