bug(//tests): Test to reproduce FP16 accuracy issue

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

tests/modules/test_fp16_compiled_modules.cpp

@@ -0,0 +1,41 @@
+#include "module_test.h"
+
+TEST_P(ModuleTests, FP16CompiledModuleIsClose) {
+    std::vector<torch::jit::IValue> jit_inputs_ivalues;
+    std::vector<torch::jit::IValue> trt_inputs_ivalues;
+    for (auto in_shape : input_shapes) {
+        auto in = at::randint(5, in_shape, {at::kCUDA});
+        in = in.to(torch::kF16);
+        jit_inputs_ivalues.push_back(in.clone());
+        trt_inputs_ivalues.push_back(in.clone());
+    }
+
+    auto extra_info = trtorch::ExtraInfo({input_shapes});
+    extra_info.op_precision = torch::kF16;
+    extra_info.strict_types = true;
+
+    auto trt_mod = trtorch::CompileGraph(mod, extra_info);
+    torch::jit::IValue trt_results_ivalues = trtorch::tests::util::RunModuleForward(trt_mod, trt_inputs_ivalues);
+    std::vector<at::Tensor> trt_results;
+    trt_results.push_back(trt_results_ivalues.toTensor());
+
+    mod.to(torch::kF16);
+    torch::jit::IValue jit_results_ivalues = trtorch::tests::util::RunModuleForward(mod, jit_inputs_ivalues);
+    std::vector<at::Tensor> jit_results;
+    jit_results.push_back(jit_results_ivalues.toTensor());
+
+    for (size_t i = 0; i < trt_results.size(); i++) {
+        ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[i], trt_results[i].reshape_as(jit_results[i]), 2e-5));
+    }
+}
+
+
+INSTANTIATE_TEST_SUITE_P(CompiledModuleForwardIsCloseSuite,
+                         ModuleTests,
+                         testing::Values(
+                            PathAndInSize({"tests/modules/resnet18.jit.pt",
+                                          {{1,3,224,224}}}),
+                            PathAndInSize({"tests/modules/resnet50.jit.pt",
+                                          {{1,3,224,224}}}),
+                            PathAndInSize({"tests/modules/mobilenet_v2.jit.pt",
+                                          {{1,3,224,224}}})));