feat: PyTorch PT2 Model Generation

This change adds the generation of PT2 format serialized PyTorch model files to the QA model generation scripts.

Author: whoisj

qa/L0_backend_python/python_based_backends/test.sh

@@ -35,6 +35,7 @@ CLIENT_LOG="./python_based_backends_client.log"
 TEST_RESULT_FILE="./test_results.txt"
 CLIENT_PY="./python_based_backends_test.py"
 GEN_PYTORCH_MODEL_PY="../../common/gen_qa_pytorch_model.py"
+GEN_PYTORCH2_MODEL_PY="../../common/gen_qa_pytorch2_model.py"
 RET=0
 
 rm -rf ${MODEL_REPOSITORY}

qa/common/gen_qa_models.py

@@ -1283,6 +1283,55 @@ def forward(self, INPUT0, INPUT1):
     traced.save(model_version_dir + "/model.pt")
 
 
+def create_libtorch2_modelfile(
+    models_dir,
+    max_batch,
+    model_version,
+    input_shape,
+    output0_shape,
+    output1_shape,
+    input_dtype,
+    output0_dtype,
+    output1_dtype,
+    swap=False,
+):
+    if not tu.validate_for_libtorch_model(
+        input_dtype,
+        output0_dtype,
+        output1_dtype,
+        input_shape,
+        output0_shape,
+        output1_shape,
+        max_batch,
+    ):
+        return
+
+    model_name = tu.get_model_name(
+        "libtorch2",
+        input_dtype,
+        output0_dtype,
+        output1_dtype,
+    )
+
+    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)
+
+    try:
+        os.makedirs(model_version_dir)
+    except OSError:
+        pass  # ignore existing dir
+
+    class AddSubNet2(nn.Module):
+        def forward(self, INPUT0, INPUT1):
+            op0 = INPUT0 + INPUT1
+            op1 = INPUT0 - INPUT1
+            return op0, op1
+
+    ep = torch.export.export(
+        AddSubNet2(), (torch.randn(*input_shape), torch.randn(*input_shape))
+    )
+    torch.export.save(ep, model_version_dir + "/model.pt2")
+
+
 def create_libtorch_modelconfig(
     models_dir,
     max_batch,
@@ -1295,6 +1344,7 @@ def create_libtorch_modelconfig(
     output1_dtype,
     output0_label_cnt,
     version_policy,
+    torch2=False,
 ):
     if not tu.validate_for_libtorch_model(
         input_dtype,
@@ -1307,6 +1357,10 @@ def create_libtorch_modelconfig(
     ):
         return
 
+    if torch2 and max_batch != 0:
+        # TorchScript models with batching are not supported in torch2 backend
+        return
+
     # Unpack version policy
     version_policy_str = "{ latest { num_versions: 1 }}"
     if version_policy is not None:
@@ -1320,15 +1374,15 @@ def create_libtorch_modelconfig(
 
     # Use a different model name for the non-batching variant
     model_name = tu.get_model_name(
-        "libtorch_nobatch" if max_batch == 0 else "libtorch",
+        "libtorch_nobatch" if max_batch == 0 else "libtorch2" if torch2 else "libtorch",
         input_dtype,
         output0_dtype,
         output1_dtype,
     )
     config_dir = models_dir + "/" + model_name
     config = """
 name: "{}"
-platform: "pytorch_libtorch"
+platform: "{}"
 max_batch_size: {}
 version_policy: {}
 input [
@@ -1358,6 +1412,7 @@ def create_libtorch_modelconfig(
 ]
 """.format(
         model_name,
+        "pytorch2_libtorch" if torch2 else "pytorch_libtorch",
         max_batch,
         version_policy_str,
         np_to_model_dtype(input_dtype),
@@ -1691,7 +1746,7 @@ def create_models(
             output1_dtype,
         )
 
-    if FLAGS.libtorch:
+    if FLAGS.libtorch or FLAGS.libtorch2:
         # max-batch 8
         create_libtorch_modelconfig(
             models_dir,
@@ -1705,6 +1760,7 @@ def create_models(
             output1_dtype,
             output0_label_cnt,
             version_policy,
+            FLAGS.libtorch2,
         )
         create_libtorch_modelfile(
             models_dir,
@@ -1716,6 +1772,7 @@ def create_models(
             input_dtype,
             output0_dtype,
             output1_dtype,
+            FLAGS.libtorch2,
         )
         # max-batch 0
         create_libtorch_modelconfig(
@@ -1730,6 +1787,7 @@ def create_models(
             output1_dtype,
             output0_label_cnt,
             version_policy,
+            FLAGS.libtorch2,
         )
         create_libtorch_modelfile(
             models_dir,
@@ -1741,6 +1799,7 @@ def create_models(
             input_dtype,
             output0_dtype,
             output1_dtype,
+            FLAGS.libtorch2,
         )
 
     if FLAGS.openvino:
@@ -1933,6 +1992,12 @@ def create_fixed_models(
         action="store_true",
         help="Generate Pytorch LibTorch models",
     )
+    parser.add_argument(
+        "--libtorch2",
+        required=False,
+        action="store_true",
+        help="Generate Pytorch LibTorch models using PT2",
+    )
     parser.add_argument(
         "--openvino",
         required=False,
@@ -1959,7 +2024,7 @@ def create_fixed_models(
         import tensorrt as trt
     if FLAGS.onnx:
         import onnx
-    if FLAGS.libtorch:
+    if FLAGS.libtorch or FLAGS.libtorch2:
         import torch
         from torch import nn
     if FLAGS.openvino:
@@ -2116,6 +2181,7 @@ def create_fixed_models(
                 create_onnx_modelfile(
                     FLAGS.models_dir, 0, 3, (16,), (16,), (16,), vt, vt, vt, swap=True
                 )
+
         if FLAGS.libtorch:
             for vt in [np.float32, np.int32, np.int16, np.int8]:
                 create_libtorch_modelfile(
@@ -2130,6 +2196,22 @@ def create_fixed_models(
                 create_libtorch_modelfile(
                     FLAGS.models_dir, 0, 3, (16,), (16,), (16,), vt, vt, vt, swap=True
                 )
+
+        if FLAGS.libtorch2:
+            for vt in [np.float32, np.int32, np.int16, np.int8]:
+                create_libtorch2_modelfile(
+                    FLAGS.models_dir, 8, 2, (16,), (16,), (16,), vt, vt, vt, swap=True
+                )
+                create_libtorch2_modelfile(
+                    FLAGS.models_dir, 8, 3, (16,), (16,), (16,), vt, vt, vt, swap=True
+                )
+                create_libtorch2_modelfile(
+                    FLAGS.models_dir, 0, 2, (16,), (16,), (16,), vt, vt, vt, swap=True
+                )
+                create_libtorch2_modelfile(
+                    FLAGS.models_dir, 0, 3, (16,), (16,), (16,), vt, vt, vt, swap=True
+                )
+
         if FLAGS.openvino:
             for vt in [np.float16, np.float32, np.int8, np.int16, np.int32]:
                 create_openvino_modelfile(

qa/common/gen_qa_pytorch2_model.py

@@ -0,0 +1,121 @@
+# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+import argparse
+import os
+
+import torch
+from torch import nn
+
+
+class AddSubNet(nn.Module):
+    def __init__(self):
+        super(AddSubNet, self).__init__()
+
+    def forward(self, input0, input1):
+        return (input0 + input1), (input0 - input1)
+
+
+def generate_model(model_dir):
+    model = AddSubNet()
+
+    os.makedirs(model_dir, exist_ok=True)
+    model_path = os.path.join(model_dir, "model.pt2")
+
+    ep = torch.export.export(model, (torch.randn(4), torch.randn(4)),)
+
+    torch.export.save(ep, model_path)
+
+
+def generate_config(config_path):
+    with open(f"{config_path}/config.pbtxt", "w") as f:
+        f.write(
+            """
+backend: "pytorch"
+input [
+  {
+    name: "INPUT0"
+    data_type: TYPE_FP32
+    dims: [ 4 ]
+  }
+]
+input [
+  {
+    name: "INPUT1"
+    data_type: TYPE_FP32
+    dims: [ 4 ]
+  }
+]
+output [
+  {
+    name: "OUTPUT0"
+    data_type: TYPE_FP32
+    dims: [ 4 ]
+  }
+]
+output [
+  {
+    name: "OUTPUT1"
+    data_type: TYPE_FP32
+    dims: [ 4 ]
+  }
+]
+"""
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-m",
+        "--model-directory",
+        type=str,
+        required=True,
+        help="The path to the model repository.",
+    )
+    parser.add_argument(
+        "--model-name",
+        type=str,
+        required=False,
+        default="add_sub_pytorch",
+        help="Model name",
+    )
+    parser.add_argument(
+        "--version",
+        type=str,
+        required=False,
+        default="1",
+        help="Model version",
+    )
+
+    args = parser.parse_args()
+
+    model_directory = os.path.join(args.model_directory, args.model_name)
+    os.makedirs(model_directory, exist_ok=True)
+
+    generate_model(model_dir=os.path.join(model_directory, args.version))
+    generate_config(model_directory)