test(serialization): Changes the serialization test to cover both

dynamic (plugin) and static, uses a pooling based model instead of
interpolate since that is handled by TensorRT now

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

Author: narendasan

tests/modules/hub.py

@@ -79,18 +79,18 @@
         script_model = torch.jit.script(m["model"])
         torch.jit.save(script_model, n + '_scripted.jit.pt')
 
-# Sample Interpolation Model (align_corners=False, for Testing Interpolate Plugin specifically)
-class Interpolate(nn.Module):
+# Sample Pool Model (for testing plugin serialization)
+class Pool(nn.Module):
     def __init__(self):
-        super(Interpolate, self).__init__()
+        super(Pool, self).__init__()
 
     def forward(self, x):
-        return F.interpolate(x, size=(10, 10, 10), align_corners=False, mode="trilinear")
+        return F.adaptive_avg_pool2d(x, (5, 5))
 
-model = Interpolate().eval().cuda()
-x = torch.ones([1, 3, 5, 5, 5]).cuda()
+model = Pool().eval().cuda()
+x = torch.ones([1, 3, 10, 10]).cuda()
 
 trace_model = torch.jit.trace(model, x)
-torch.jit.save(trace_model, "interpolate_traced.jit.pt")
+torch.jit.save(trace_model, "pooling_traced.jit.pt")
 
 

tests/modules/test_serialization.cpp

@@ -1,5 +1,23 @@
 #include "module_test.h"
 
+std::vector<trtorch::ExtraInfo::InputRange> toInputRangesDynamic(std::vector<std::vector<int64_t>> opts) {
+    std::vector<trtorch::ExtraInfo::InputRange> a;
+
+    for (auto opt : opts) {
+        std::vector<int64_t> min_range(opt);
+        std::vector<int64_t> max_range(opt);
+
+        min_range[3] = ceil(opt[3]/2.0);
+        max_range[3] = 2*opt[3];
+        min_range[2] = ceil(opt[2]/2.0);
+        max_range[2] = 2*opt[2];
+
+        a.push_back(trtorch::ExtraInfo::InputRange(min_range, opt, max_range));
+    }
+
+    return std::move(a);
+}
+
 TEST_P(ModuleTests, SerializedModuleIsStillCorrect) {
     std::vector<torch::jit::IValue> post_serialized_inputs_ivalues;
     std::vector<torch::jit::IValue> pre_serialized_inputs_ivalues;
@@ -26,11 +44,37 @@ TEST_P(ModuleTests, SerializedModuleIsStillCorrect) {
     }
 }
 
+TEST_P(ModuleTests, SerializedDynamicModuleIsStillCorrect) {
+    std::vector<torch::jit::IValue> post_serialized_inputs_ivalues;
+    std::vector<torch::jit::IValue> pre_serialized_inputs_ivalues;
+    for (auto in_shape : input_shapes) {
+        auto in = at::randint(5, in_shape, {at::kCUDA});
+        post_serialized_inputs_ivalues.push_back(in.clone());
+        pre_serialized_inputs_ivalues.push_back(in.clone());
+    }
+
+    auto pre_serialized_mod = trtorch::CompileGraph(mod, toInputRangesDynamic(input_shapes));
+    torch::jit::IValue pre_serialized_results_ivalues = trtorch::tests::util::RunModuleForward(pre_serialized_mod, pre_serialized_inputs_ivalues);
+    std::vector<at::Tensor> pre_serialized_results;
+    pre_serialized_results.push_back(pre_serialized_results_ivalues.toTensor());
+
+    pre_serialized_mod.save("test_serialization_mod.ts");
+    auto post_serialized_mod = torch::jit::load("test_serialization_mod.ts");
+
+    torch::jit::IValue post_serialized_results_ivalues = trtorch::tests::util::RunModuleForward(post_serialized_mod, post_serialized_inputs_ivalues);
+    std::vector<at::Tensor> post_serialized_results;
+    post_serialized_results.push_back(post_serialized_results_ivalues.toTensor());
+
+    for (size_t i = 0; i < pre_serialized_results.size(); i++) {
+        ASSERT_TRUE(trtorch::tests::util::almostEqual(post_serialized_results[i], pre_serialized_results[i].reshape_as(post_serialized_results[i]), 2e-5));
+    }
+}
+
 
 INSTANTIATE_TEST_SUITE_P(CompiledModuleForwardIsCloseSuite,
                          ModuleTests,
                          testing::Values(
                             PathAndInSize({"tests/modules/resnet18_traced.jit.pt",
                                           {{1,3,224,224}}}),
-                            PathAndInSize({"tests/modules/interpolate_traced.jit.pt",
-                                          {{1,3,5,5,5}}})));
+                            PathAndInSize({"tests/modules/pooling_traced.jit.pt",
+                                          {{1,3,10,10}}})));