Multiple Signature support in Mobile

Mobile
1

Is there a way to compile a Module with multiple forward-like functions for pytorch mobile?

Imagine you have a CNN trunk model for image analysis with multiple heads for different tasks, but you only ever need one task at a time – so you might have segment(), recognize(), validate() methods that share 90% of the logic they’re using under the hood. Do I have to compile each of these as separate models (resulting in a lot of duplication) or can I save the module once and have access to all three paths?

If on-device training is on the roadmap I assume this is too, but I haven’t been able to find anything in the docs about it.

Thanks!

2

Yes. To compile a method other than forward (and recursively compile anything it calls), add the @torch.jit.export decorator to the method.

See torch.jit.script — PyTorch 1.9.0 documentation

1 Like
3

Thank you! Coming from tensorflow/tflite I did not expect it to be so straightforward…

1 Like