[GitHub] [tvm] trevor-m commented on a change in pull request #6955: Dynamic Batch Support for TRT

[GitBox <gi...@apache.org>]

trevor-m commented on a change in pull request #6955:
URL: https://github.com/apache/tvm/pull/6955#discussion_r532759880



##########
File path: python/tvm/relay/op/contrib/tensorrt.py
##########
@@ -607,8 +650,9 @@ def reshape_annotate_fn(expr):  # pylint: disable=unused-variable
         for i, value in enumerate(new_shape):
             if value == -1:
                 new_shape[i] = original_volume // np.prod([x for x in new_shape if x != -1])
+        # Remove batch dimension and see if volumes match
         if shape[0] != new_shape[0]:
-            logger.info("reshape: can't modify batch dimension.")
+            print("reshape: can't modify batch dimension.")

Review comment:
       Change back to `logger.info`

##########
File path: python/tvm/relay/op/contrib/tensorrt.py
##########
@@ -808,18 +884,22 @@ def is_valid_subgraph(params, body):
                     if len(tupe_type.shape) == 0:
                         logger.info("tensorrt: scalar inputs not supported")
                         return False
-                    input_batch_sizes.append(int(tupe_type.shape[0]))
+                    if not isinstance(tupe_type.shape[0], tvm.tir.expr.Any):
+                        input_batch_sizes.append(int(tupe_type.shape[0]))
             else:
                 # Scalar inputs not allowed
                 if len(var.checked_type.shape) == 0:
                     logger.info("tensorrt: scalar inputs not supported")
                     return False
-                input_batch_sizes.append(int(var.checked_type.shape[0]))
+                if not isinstance(var.checked_type.shape[0], tvm.tir.expr.Any):

Review comment:
       By not checking the batch size here (and above for tuple type input), we open ourselves up to a possibility where we have partitioned a subgraph for TRT which cannot be executed if the inputs have different batch sizes.
   
   Maybe we should only allow subgraphs with dynamic batch sizes if there is only one input.
   
   This limitation only applies to implicit batch mode in TRT.

##########
File path: src/runtime/contrib/tensorrt/tensorrt_runtime.cc
##########
@@ -240,30 +249,35 @@ class TensorRTRuntime : public JSONRuntimeBase {
     helper.DeclareField("inputs", &engine_and_context.inputs);
     helper.DeclareField("outputs", &engine_and_context.outputs);
     helper.ReadAllFields(&reader);
-    trt_engine_cache_[symbol_name_] = engine_and_context;
+    const int batch_size = 1;

Review comment:
       We should probably add the batch size of the engine in the key (file name) of the cached engine.

##########
File path: tests/python/contrib/test_tensorrt.py
##########
@@ -1034,5 +1039,119 @@ def set_func_attr(func, compile_name, symbol_name):
     tvm.ir.assert_structural_equal(mod_trt, mod_exp, map_free_vars=True)
 
 
+def convert_traced_model_to_vm_trt(

Review comment:
       I don't think we should have these pytorch tests in the TRT tests.

##########
File path: tests/python/contrib/test_tensorrt.py
##########
@@ -1034,5 +1039,119 @@ def set_func_attr(func, compile_name, symbol_name):
     tvm.ir.assert_structural_equal(mod_trt, mod_exp, map_free_vars=True)
 
 
+def convert_traced_model_to_vm_trt(
+    traced_module: torch.jit.TopLevelTracedModule, np_sample_input: np.ndarray, target: str
+) -> tvm.runtime.vm.Executable:
+    """
+    This function converts a traced pytorch model to VM + TRT.
+    """
+    input_shape = np_sample_input.shape
+    input_name = "input0"
+    shape_list = [(input_name, input_shape)]
+    mod, params = relay.frontend.from_pytorch(traced_module, shape_list)
+    mod, config = tensorrt.partition_for_tensorrt(mod, params, remove_no_mac_subgraphs=True)
+    with tvm.transform.PassContext(opt_level=3, disabled_pass=["FoldScaleAxis"]):
+        vm_trt_exec = relay.vm.compile(mod, target=target, params=params)
+
+    return vm_trt_exec
+
+
+def test_maskrcnn_resnet50() -> None:

Review comment:
       What properties of this particular model are we trying to test here? Is it really specific to TRT integrattion? Can we make a simple test case which replicates the same behavior instead - i.e. single conv2d op with dynamic batch size in TRT.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org