This sounds like in theory and makes you curse in practice. Trained TF model -> deployment using TF Lite is the most robust pipeline currently as far as I know.
There are lots of improvements going into pytorch for mobile at the moment, but for the moment I'll wait and see how it turns out - I didn't have much fun with caffe2 when "train in pytorch and deploy with caffe2" was the storyline FB pushed (e.g. problems with binary size and slow depthwise convolutions) so not too eager to migrate back at the moment.
It's difficult because the graph is defined running Python code, which can basically do anything. As I understand it you have to actually run the model and sort of introspect it to work out what the ONNX graph should be.
Whereas the Tensorflow API actually creates a static graph that can be easily converted to ONNX.
Edit: They talk about this problem in the article:
> Although straightforward, tracing has its downsides. For example, it can’t capture control flow that didn’t execute. For example, it can’t capture the false block of a conditional if it executed the true block.
> Script mode takes a function/class, reinterprets the Python code and directly outputs the TorchScript IR. This allows it to support arbitrary code, however it essentially needs to reinterpret Python.
I've done this with OpenVino and it works well enough. There are some gotchas with Onnx, some layers break or have weird bugs (upsample in particular used to be a problem).
There are lots of improvements going into pytorch for mobile at the moment, but for the moment I'll wait and see how it turns out - I didn't have much fun with caffe2 when "train in pytorch and deploy with caffe2" was the storyline FB pushed (e.g. problems with binary size and slow depthwise convolutions) so not too eager to migrate back at the moment.