Can Keras be used to painlessly convert PyTorch/Jax models to TF models ?

[BenCrulis]

I recently had to convert PyTorch models to Tensorflow models and then to TFLite models for export to Android devices.

Existing options are somewhat cumbersome or limited, I encountered issues when trying to use ONNX as a bridge, and the Nobuco library what difficult to debug to make it work, even though I managed to make the pipeline work with Nobuco in the end.

Now that Keras is multi-backend, is there an easy way to use Keras as a bridge for converting between native PyTorch and Jax models to Tensorflow models ?
And if there is currently no way, is it possible to imagine writing a small library based on Keras that does this ?

[abhaskumarsinha]

You can use Keras to convert your model in PyTorch into Keras manually. Once the Keras code has been written, it becomes very easy to run that Keras code into any of the three backends - TensorFlow, PyTorch or JAX.

Keras is not an automated converter tool, it acts like a framework that is more backend-independent.

For simple models, writing such converters is easy, but for complex ones, dynamic graphing, manual intervention is a must.

Best Regards,
Abhas

[BenCrulis]

So if I understand correctly, there is currently no way of (1) automatically export a PyTorch nn.Module object to Keras and then (2) export the resulting Keras model to Tensorflow ?

I imagine (2) is implemented somewhere or can be easily implemented, but regarding (1), do we absolutely need to copy paste the PyTorch code and completely re-write it to Keras even in simple cases ?

I see in the Keras doc that there is a TorchModuleWrapper class, is it too limited for this ?

[abhaskumarsinha]

1. Writing a tool that does converts PyTorch nn.Module to Keras shouldn't be hard, but since Keras 3 is fairly new there isn't any such tool for now.
2. Keras Model can be exported to TensorFlow (weights ONLY). To export the code, a little modification might be required.

To be honest, I haven't used TorchModuleWrapper before, you'd have to test a bit to check if that supports custom Torch layers or not. But again that'd not port or convert any of your code to TensorFlow or PyTorch, it would rather act as an emulator.

[fchollet]

TorchModuleWrapper is for using arbitrary torch modules as Keras layers when using the torch backend. It cannot run with TF or JAX.

[BAW2501]

As of now TTBOMK, no utility converts the actual model object or code to and from Keras ( Keras 3).

manual conversion is in order( shouldn't be too tedious given you're planning to run the models on edge devices) ,

as for weight conversion for keras( Keras 2) ONNX is somewhat viable in my experience
torch.onnx.export -> onnx.load -> onnx_to_keras

[pass-lin]

no problem. my lib is written by pure Keras API. It can easily convert torch model to jax or tf
https://github.com/pass-lin/bert4keras3

[abhaskumarsinha]

Thank you for enlightning us about that. Can you tell if that ONLY works for BERT or vision models too?

[pass-lin]

Thank you for enlightning us about that. Can you tell if that ONLY works for BERT or vision models too?
my lib also support LLM(gemma,llama,qwen) ,T5 and other bert like model.
but I am sorry that main support Chinese model

[abhaskumarsinha]

Okay, I got it. No worries. Thank You.