I was looking for a new option for converting PyTorch models to TensorFlow.js (since my previous method is no longer maintained), and I came across this project for translating PyTorch models directly to TensorFlow:
- nobuco: Pytorch to Tensorflow conversion made intuitive.
The project converts a PyTorch model to Keras, which you can then save in TensorFlow’s SavedModel format for use with the TensorFlow.js converter tool.
I used the project to convert a YOLOX Tiny model from PyTorch to TensorFlow.js and found the resulting model to be about twice as fast compared to my previous PyTorch→ONNX→TensorFlow→TensorFlow.js approach.
You can compare the inference speed of models exported using the two approaches with the in-browser demos linked below:
The one downside is I have not been able to get dynamic input dimensions working for the YOLOX model exported using Nobuco’s default conversion. Nobuco lets you customize how PyTorch models get converted, but I have not attempted to get dynamic input dimensions working with that approach.
Here is the Jupyter notebook I used to convert the YOLOX model with nobuco: