DNN Model Splitting through TVM

I started looking into TVM a few days ago and it looks really exiting and might be suitable for an experiment I plan to run. My understanding of the compilation process is still quite limited so this might be a very basic question.

To my current understanding I can specify a model in TF or Pytorch and then have it compiled into something that runs optimized for the target architecture and environment.

Now to my question: Would it be possible for me to extend TVM (eg. through customized passes) to split a network at a certain layer and obtain two independent networks deployable independently on the TVM runtime? In this scenario the output of the first network would be fed as the input to the second one where both networks do not run on the same TVM runtime.

Concretly I was wondering whether for an image classification network we can split the network into two networks where the first network portion would then be a general image encoder and the second portion the task-specific section.

If you have any idea how to achieve this inside the TVM environment I would be grateful for your insights.

Thanks,
Daniel

1 Like