Does tvm support dynamic input shape?

I have a LSTM model written with pytorch, and first i convert it to onnx model, this model has a dynamic input shape represent as: [batch_size, seq_number], so when i compile this model with: relay.frontend.from_onnx(onnx_model), there will convert the dynamic shape with type Any . so when execute at ./relay/frontend/onnx.py: X_steps = unbind(X, axis=0)./relay/frontend/common.py: _op.split(data, selections, axis) , there raise exception: Any type has no len attributes. so, does tvm not support dynamic input shape for lstm?

2 Likes

Can anybody reply this question? thx so much.

In principle, TVM can compile models with dynamic input shape. But currently performance would be terrible due to the lack of dynamic shape tuning. Moreover, many of conversion functions in the importers make an assumption that shapes are not dynamic (such as the error you found). So support for compiling dynamic input shape is extremely poor.

thanks for you rapid reply, because most Lstm based model, will required dynamic input sequence, so this will limit the capacity of tvm. are there some suggestion for process this lstm model, or is there a plan to extend dynamic input for lstm? thx so much~

Hi, is there any update about the dynamic input shape support now? Thanks!

Hi @puddingfjz, we are working on Relax(Relay Next), which has first-class symbolic shape support. Here is a simple dynamic input shape example. Please check out the overview of Relax and our TVMCon talk for more information.

@ziheng has worked on dynamic-shape workload tuning: DietCode: Automatic Optimization for Dynamic Tensor Programs - TVMCon 2021.

@yongwww is working on supporting LSTM on Relax, and adding control flow support(e.g., closure, while_loop) to Relax.

4 Likes

Thanks for your quick reply!

Is the Dietcode work integrated into TVM now? Are there any examples to show how to use tvm to optimize dynamic input shape operators?

Is there examples of supporting LSTM or which part of code is related to this?

Dietcode is not upstreamed yet afaik, cc @ziheng, @ArmageddonKnight.

We will add an LSTM demo with dynamic input shape to the Relax repo probably next month, we will keep you updated! :slight_smile:

Good summary, thanks. Hope dynamic feature is supported well in TVM ASAP.

Does the demo ready now?

Hi @gfvvz, @yongwww is working on Relax control flow support which is necessary to support LSTM. After it, we will also need to support TensorArray. We have pinned down the design after discussing with the community in the Relax open dev meetings, now Yong is working on the implementations.

We estimate to deliver the LSTM demo(task list and tracking) with dynamic input shape in June. Stay tuned!

Is there any updates regarding the models with dynamic shape inputs?

I run a model with dynamic inputs and there wan’t any improvement for the inference time.

Hi @fnakhaee, did you run the model with dynamic inputs in Relay or Relax?

Currently in Relax, we are working on matching Relay performance for static workloads by adding passes like FuseOps, after that we will shift our focus to improving performance for dynamic workloads, which involves works around dynamic shape tuning as masahi mentioned before and Relax passes such as dynamic memory planning.

Feel free to check out Relax pre-RFC for more details and the discussion we had in the TVM community meeting yesterday: https://youtu.be/2aYWGOYmDFY.

Hi, yuchen. I saw your code example abount dynamic shape.

I changed the target and virtual machine to cuda.

But it got a error about forgot binding thread, so does current relax support dynamic shape with gpu. Or, I need to add other code. @yuchenj

Hi @chenugray,

The default TIR PrimFunc emitted by EmitTE do not have thread binding construct, so Relax currently relied on MetaSchedule to do thread binding when targeting GPU, you just need to perform one tuning trial per task (subgraph) when targeting GPU.

You can find a test case here: https://github.com/tlc-pack/relax/blob/relax/tests/python/relax/test_autotir_integration.py#L131.

@yuchenj ,here is my code example. I wanna run dynamic batch of customer model on gpu.

It seems like MetaSchedule can’t tune on dynmiac tasks.

from __future__ import annotations  # must import to defer parsing of annotations
import numpy as np
import tvm
from tvm import relax, tir
from tvm.relax.testing import nn
import tempfile
from tvm import meta_schedule as ms

builder = relax.BlockBuilder()

input_size = 784
hidden_sizes = [128, 64]
output_size = 10
target = tvm.target.Target(
    "cuda --host=llvm --max_threads_per_block=1024 --max_shared_memory_per_block=49152")
database = ms.database.MemoryDatabase()

with builder.function(name="main"):
    model = nn.Sequential(
        nn.Linear(input_size, hidden_sizes[0]),
        nn.ReLU(),
        nn.Linear(hidden_sizes[0], hidden_sizes[1]),
        nn.ReLU(),
        nn.Linear(hidden_sizes[1], output_size),
        nn.LogSoftmax(),
    )
    # n is a symbolic variable to represent a dynamic batch size
    n = tir.Var("n", "int64")
    data = nn.Placeholder((n, input_size), name="data")
    output = model(data)
    params = [data] + model.parameters()
    builder.emit_func_output(output, params=params)

mod = builder.get()
with tempfile.TemporaryDirectory() as work_dir:
    db = ms.relax_integration.tune_relax(
        mod=mod,
        target=target,
        params=None,
        num_trials_per_iter=2,
        max_trials_per_task=4,
        max_trials_global=4,
        work_dir=work_dir,
        database=database,
    )

relax_ex = ms.relax_integration.compile_relax(
    db, mod=mod, target=target, params=None)
vm = relax.VirtualMachine(relax_ex, tvm.cuda())
params = nn.init_params(mod)
# init parameters
# the input data has a minibatch size of 3
data = tvm.nd.array(np.random.rand(3, input_size).astype(np.float32))
res = vm["main"](data, *params)

I got error like below:

2: tvm::tir::StmtFunctor<tvm::tir::TResult (tvm::tir::Stmt const&)>::InitVTable()::{lambda(tvm::runtime::ObjectRef const&, tvm::tir::StmtFunctor<tvm::tir::TResult (tvm::tir::Stmt const&)>*)#4}::_FUN(tvm::runtime::ObjectRef const&, tvm::tir::StmtFunctor<tvm::tir::TResult (tvm::tir::Stmt const&)>*)
        at /home/lei.zhang/src_code/cmccperf/relax/include/tvm/tir/stmt_functor.h:115
  1: tvm::tir::StmtFunctor<tvm::tir::TResult (tvm::tir::Stmt const&)>::InitVTable()::{lambda(tvm::runtime::ObjectRef const&, tvm::tir::StmtFunctor<tvm::tir::TResult (tvm::tir::Stmt const&)>*)#4}::operator()(tvm::runtime::ObjectRef const&, tvm::tir::StmtFunctor<tvm::tir::TResult (tvm::tir::Stmt const&)>*) const
        at /home/lei.zhang/src_code/cmccperf/relax/include/tvm/tir/stmt_functor.h:115
  0: tvm::tir::FlopEstimator::VisitStmt_(tvm::tir::ForNode const*)
        at /home/lei.zhang/src_code/cmccperf/relax/src/tir/analysis/estimate_flops.cc:143
  File "/home/lei.zhang/src_code/cmccperf/relax/src/tir/analysis/estimate_flops.cc", line 143
TVMError: 
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------
  Check failed: (int_imm) is false: TypeError: Expect the extent of a loop to be IntImm, but gets: tir.Var

Does that need dietcode to tune task for dynamic? And when will it be integrated in TVM?

Hi @chenugray,

The current meta-schedule does not support tuning dynamic workloads. I don’t know the plan of dietcode, but dynamic tuning of course is an important part of Relax, we will go for it (but no clear timeline for now either)

Hi @Hzfengsy ,

I am also interested in dynamic shape feature in TVM and have tried several Relax demo for dynamic input shape, such as nn_module.py and all passed. Meanwhile I want to load a TorchScript model, then convert it to Relax IRModule ONLY once . Then use different input shape as its input data. Is there any toy demo? Thanks~

Unfortunately, there are no vertical examples right now. We are working towards it

Excited to hear that and hope the example can be released ASAP.