refactor: Refactor testing to use cosine similarity, remove redundancy models and restructuring

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

tests/core/lowering/test_module_fallback_passes.cpp

@@ -124,5 +124,5 @@ TEST(Lowering, LowerAndPartitionSimpleModuleFallbackCorrectly) {
   }
 
   auto trt_results = trt_mod.forward(trt_inputs_ivalues).toTensor();
-  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results, trt_results, 2e-6));
+  ASSERT_TRUE(torch_tensorrt::tests::util::cosineSimEqual(jit_results, trt_results, 0.99));
 }

tests/core/partitioning/BUILD

@@ -55,21 +55,6 @@ cc_test(
     }),
 )
 
-cc_test(
-    name = "test_fallback_graph_output",
-    srcs = ["test_fallback_graph_output.cpp"],
-    data = [
-        ":jit_models",
-    ],
-    deps = [
-        "//tests/util",
-        "@googletest//:gtest_main",
-    ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch//:libtorch"],
-    }),
-)
-
 cc_test(
     name = "test_loop_fallback",
     srcs = ["test_loop_fallback.cpp"],
@@ -104,7 +89,6 @@ test_suite(
     name = "partitioning_tests",
     tests = [
         ":test_conditionals",
-        ":test_fallback_graph_output",
         ":test_loading_model",
         ":test_loop_fallback",
         ":test_resolve_nontensor_inputs",

tests/core/partitioning/test_fallback_graph_output.cpp

@@ -13,12 +13,9 @@ test_suite(
     name = "api_tests",
     tests = [
         ":test_collections",
-        ":test_compiled_modules",
         ":test_default_input_types",
         ":test_example_tensors",
-        ":test_module_fallback",
         ":test_modules_as_engines",
-        ":test_multiple_registered_engines",
         ":test_runtime_thread_safety",
         ":test_serialization",
     ],
@@ -28,12 +25,9 @@ test_suite(
     name = "aarch64_api_tests",
     tests = [
         ":test_collections",
-        ":test_compiled_modules",
         ":test_default_input_types",
         ":test_example_tensors",
-        ":test_module_fallback",
         ":test_modules_as_engines",
-        ":test_multiple_registered_engines",
         ":test_runtime_thread_safety",
         ":test_serialization",
     ],
@@ -72,21 +66,6 @@ cc_test(
     ],
 )
 
-cc_test(
-    name = "test_multiple_registered_engines",
-    srcs = ["test_multiple_registered_engines.cpp"],
-    data = [
-        "//tests/modules:jit_models",
-    ],
-    deps = [
-        "//tests/util",
-        "@googletest//:gtest_main",
-    ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch//:libtorch"],
-    }),
-)
-
 cc_test(
     name = "test_modules_as_engines",
     timeout = "long",
@@ -110,21 +89,6 @@ cc_test(
     ],
 )
 
-cc_test(
-    name = "test_module_fallback",
-    srcs = ["test_module_fallback.cpp"],
-    data = [
-        "//tests/modules:jit_models",
-    ],
-    deps = [
-        "//tests/util",
-        "@googletest//:gtest_main",
-    ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch//:libtorch"],
-    }),
-)
-
 cc_test(
     name = "test_collections",
     srcs = ["test_collections.cpp"],
@@ -140,17 +104,6 @@ cc_test(
     }),
 )
 
-cc_test(
-    name = "test_compiled_modules",
-    srcs = ["test_compiled_modules.cpp"],
-    data = [
-        "//tests/modules:jit_models",
-    ],
-    deps = [
-        ":cpp_api_test",
-    ],
-)
-
 cc_test(
     name = "test_multi_gpu_serde",
     srcs = ["test_multi_gpu_serde.cpp"],

tests/cpp/BUILD

@@ -15,40 +15,7 @@ TEST_P(CppAPITests, ModuleAsEngineIsClose) {
   auto trt_results = torch_tensorrt::tests::util::RunModuleForwardAsEngine(mod, inputs);
 
   ASSERT_TRUE(
-      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), threshold));
-}
-
-TEST_P(CppAPITests, ModuleToEngineToModuleIsClose) {
-  std::vector<at::Tensor> inputs;
-  std::vector<torch::jit::IValue> inputs_ivalues;
-  for (uint64_t i = 0; i < input_shapes.size(); i++) {
-    inputs.push_back(at::randint(5, input_shapes[i], {at::kCUDA}).to(input_types[i]));
-    inputs_ivalues.push_back(inputs[inputs.size() - 1].clone());
-  }
-
-  torch::jit::IValue jit_results_ivalues = torch_tensorrt::tests::util::RunModuleForward(mod, inputs_ivalues);
-  std::vector<at::Tensor> jit_results;
-  jit_results.push_back(jit_results_ivalues.toTensor());
-
-  std::vector<c10::ArrayRef<int64_t>> input_ranges;
-  for (auto in : inputs) {
-    input_ranges.push_back(in.sizes());
-  }
-
-  auto compile_spec = torch_tensorrt::ts::CompileSpec({input_ranges});
-  int device_id = 0;
-  cudaGetDevice(&device_id);
-  compile_spec.device.device_type = torch_tensorrt::Device::DeviceType::kGPU;
-  compile_spec.device.gpu_id = device_id;
-  auto engine = torch_tensorrt::ts::convert_method_to_trt_engine(mod, "forward", input_ranges);
-  auto trt_mod = torch_tensorrt::ts::embed_engine_in_new_module(engine, compile_spec.device);
-
-  torch::jit::IValue trt_results_ivalues = torch_tensorrt::tests::util::RunModuleForward(trt_mod, inputs_ivalues);
-  std::vector<at::Tensor> trt_results;
-  trt_results.push_back(trt_results_ivalues.toTensor());
-
-  ASSERT_TRUE(
-      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), threshold));
+      torch_tensorrt::tests::util::cosineSimEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), threshold));
 }
 
 #ifndef DISABLE_TEST_IN_CI
@@ -57,12 +24,8 @@ INSTANTIATE_TEST_SUITE_P(
     ModuleAsEngineForwardIsCloseSuite,
     CppAPITests,
     testing::Values(
-        PathAndInput({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
-        PathAndInput({"tests/modules/resnet50_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
-        PathAndInput({"tests/modules/mobilenet_v2_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
-        PathAndInput({"tests/modules/resnet18_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
-        PathAndInput({"tests/modules/resnet50_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
-        PathAndInput({"tests/modules/mobilenet_v2_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
-        PathAndInput({"tests/modules/efficientnet_b0_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 1e-4}),
-        PathAndInput({"tests/modules/vit_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 8e-2})));
+        PathAndInput({"tests/modules/resnet50_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 0.99}),
+        PathAndInput({"tests/modules/mobilenet_v2_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 0.99}),
+        PathAndInput({"tests/modules/efficientnet_b0_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 0.99}),
+        PathAndInput({"tests/modules/vit_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 0.99})));
 #endif