Dynamic input/output shapes (batch size)

Questions
#1

Hello everyone,

I am currently working on a project where I need to handle dynamic shapes (in my case dynamic batch sizes) with a ONNX model. I saw in mid-2021 that Auto Scheduler didn’t handle Relay.Any() and future work needed to be done. The workaround I chose is optimizing the model after fixing the batch size, and this for multiple batch sizes, and then loading the corresponding exported library (.so file) like this everytime I need to switch batch size. However, it takes a lot of time (as tuning one model can take multiple hours). Furthermore, I have memory (RAM and disk) limitations so the workaround is a bit problematic.

Are there another way to handle dynamic shapes or do dynamic shapes be supported soon? Thank you!

2 Likes
#2

Yes. @yuchenj is working on the next-generational high-level IR called Relax that is able to represent dynamic shapes properly, and Bojian Zheng (UofT) and @ziheng are working on dynamic shape tuning called DietCode

1 Like
#3

Okay, thank you very much for the quick answer! Is there an estimated release date already for these features? Just to know.

Thank you !

#4

Hi @Upscale4152, currently we are doing open development for Relax in a separate repo: Relax repo. We will release a Relax Roadmap RFC this week as a start to discuss the upstreaming plan with the broader TVM community.

This Jupyter notebook contains a set of demos showing how Relax handles dynamic-shape workloads, feel free to try it out! For more details about Relax, please check out the Relax pre-RFC, the TVM Community Meeting from 04/20/2022 here, the TVMCon 2021 talk here, and design docs here.

4 Likes
#5

Okay that sounds great, I will take a look !

Thank for all of the materials !

#6

@yuchenj Is it possible to dump the C source from relax? like we normally do from relay.build()'s lib.get_source().

#7

@harishch4, you can get the C source for the compiled kernel library by:

exec = relax.vm.build(relax_mod: IRModule, target)
c_str = exec.mod.imported_modules[0].get_source()
#8

what is the status now? can TVM suport dynamic batch size with auto turing?