Repository: pytorch/TensorRT

Issue:
Create some sort of serialization / deserialization functionality
With INT8 about to land, would be a pain to have to calibrate from scratch every time. There should be some mechanism to save and load modules with the TRT engine included. 


Discussion:
That would be awesome to have proper serialization! I would also add that for large scale deployments this is a must have even for fp16. As if you have a lot of models, converting all of them on restart would be quite expensive. And in general in some systems doing tuning in production is generally undesirable. 

Resolving patch:
diff --git a/py/trtorch/__init__.py b/py/trtorch/__init__.py
--- a/py/trtorch/__init__.py
+++ b/py/trtorch/__init__.py
@@ -11,3 +11,7 @@
 from trtorch._compiler import *
 from trtorch._types import *
 from trtorch import logging
+
+def _register_with_torch():
+    trtorch_dir = os.path.dirname(__file__)
+    torch.ops.load_library(trtorch_dir + '/lib/trtorch.so')
\ No newline at end of file


Test patch:
diff --git a/tests/modules/BUILD b/tests/modules/BUILD
--- a/tests/modules/BUILD
+++ b/tests/modules/BUILD
@@ -1,3 +1,5 @@
+package(default_visibility = ["//visibility:public"])
+
 config_setting(
     name = "use_pre_cxx11_abi",
     values = {
@@ -15,7 +17,19 @@ test_suite(
     tests = [
         ":test_modules_as_engines",
         ":test_compiled_modules",
-        ":test_multiple_registered_engines"
+        ":test_multiple_registered_engines",
+        ":test_serialization"
+    ]
+)
+
+cc_test(
+    name = "test_serialization",
+    srcs = ["test_serialization.cpp"],
+    deps = [
+        ":module_test",
+    ],
+    data = [
+        ":jit_models"
     ]
 )
 
diff --git a/tests/modules/hub.py b/tests/modules/hub.py
--- a/tests/modules/hub.py
+++ b/tests/modules/hub.py
@@ -2,64 +2,68 @@
 import torchvision.models as models
 
 models = {
-          "alexnet": {
-            "model": models.alexnet(pretrained=True),
-            "path": "both"
-          },
-          "vgg16": {
-            "model": models.vgg16(pretrained=True),
-            "path": "both"
-          },
-          "squeezenet": {
-            "model": models.squeezenet1_0(pretrained=True),
-            "path": "both"
-          },
-          "densenet": {
-            "model": models.densenet161(pretrained=True),
-            "path": "both"
-          },
-          "inception_v3": {
-            "model": models.inception_v3(pretrained=True),
-            "path": "both"
-          },
-          #"googlenet": models.googlenet(pretrained=True),
-          "shufflenet": {
-            "model": models.shufflenet_v2_x1_0(pretrained=True),
-            "path": "both"
-          },
-          "mobilenet_v2": {
-            "model": models.mobilenet_v2(pretrained=True),
-            "path": "both"
-          },
-          "resnext50_32x4d": {
-            "model":  models.resnext50_32x4d(pretrained=True),
-            "path": "both"
-          },
-          "wideresnet50_2": {
-            "model": models.wide_resnet50_2(pretrained=True),
-            "path": "both"
-          },
-          "mnasnet": {
-            "model": models.mnasnet1_0(pretrained=True),
-            "path": "both"
-          },
-          "resnet18": {
-            "model": torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True),
-            "path": "both"
-          },
-          "resnet50": {
-            "model":torch.hub.load('pytorch/vision:v0.6.0', 'resnet50', pretrained=True),
-            "path": "both"
-          },
-          "fcn_resnet101": {
-            "model": torch.hub.load('pytorch/vision:v0.6.0', 'fcn_resnet101', pretrained=True),
-            "path": "script"
-          },
-          "ssd": {
-                "model": torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd', model_math="fp32"),
-                "path": "trace"
-          }
+    "alexnet": {
+        "model": models.alexnet(pretrained=True),
+        "path": "both"
+    },
+    "vgg16": {
+        "model": models.vgg16(pretrained=True),
+        "path": "both"
+    },
+    "squeezenet": {
+        "model": models.squeezenet1_0(pretrained=True),
+        "path": "both"
+    },
+    "densenet": {
+        "model": models.densenet161(pretrained=True),
+        "path": "both"
+    },
+    "inception_v3": {
+        "model": models.inception_v3(pretrained=True),
+        "path": "both"
+    },
+    #"googlenet": models.googlenet(pretrained=True),
+    "shufflenet": {
+        "model": models.shufflenet_v2_x1_0(pretrained=True),
+        "path": "both"
+    },
+    "mobilenet_v2": {
+        "model": models.mobilenet_v2(pretrained=True),
+        "path": "both"
+    },
+    "resnext50_32x4d": {
+        "model":  models.resnext50_32x4d(pretrained=True),
+        "path": "both"
+    },
+    "wideresnet50_2": {
+        "model": models.wide_resnet50_2(pretrained=True),
+        "path": "both"
+    },
+    "mnasnet": {
+        "model": models.mnasnet1_0(pretrained=True),
+        "path": "both"
+    },
+    "resnet18": {
+        "model": torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True),
+        "path": "both"
+    },
+    "resnet50": {
+        "model":torch.hub.load('pytorch/vision:v0.6.0', 'resnet50', pretrained=True),
+        "path": "both"
+    },
+    "fcn_resnet101": {
+        "model": torch.hub.load('pytorch/vision:v0.6.0', 'fcn_resnet101', pretrained=True),
+        "path": "script"
+    },
+    "ssd": {
+        "model": torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd', model_math="fp32"),
+        "path": "trace"
+    },
+    "faster_rcnn": {
+        "model": models.detection.fasterrcnn_resnet50_fpn(pretrained=True),
+        "path": "script"
     }
+}
 
 for n, m in models.items():
     print("Downloading {}".format(n))
diff --git a/tests/modules/test_serialization.cpp b/tests/modules/test_serialization.cpp
new file mode 100644
--- /dev/null
+++ b/tests/modules/test_serialization.cpp
@@ -0,0 +1,34 @@
+#include "module_test.h"
+
+TEST_P(ModuleTests, SerializedModuleIsStillCorrect) {
+    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, 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}}})));