[Python] Convert custom YOLOX model (.pth) to tflite model

I have trained YOLOX using this repo on a custom dataset with 8 classes. The final output is a .pth file which I need to convert to a .tflite format but I am completely stuck.

What I tried so far:

1. Use ai-edge-torch but I can't install the required library ai-edge-litert-nightly because it's only available on Linux and I'm using Windows.
2. Use the official Collab the same repo provides, but it crashes at the import ai_edge_torch line. (Same is true for any other Collab that uses ai-edge-troch) The error reads:

/usr/local/lib/python3.11/dist-packages/tensorflow/python/framework/load_library.py in load_library(library_location)
149
150 for lib in kernel_libraries:
--> 151 py_tf.TF_LoadLibrary(lib)
152
153 else:

NotFoundError: /usr/local/lib/python3.11/dist-packages/tensorflow/core/kernels/libtfkernel_sobol_op.so: undefined symbol: _ZN10tensorflow15TensorShapeBaseINS_11TensorShapeEEC2EN4absl12lts_202308024SpanIKlEE

1. Via onnx - I tried to follow this tutorial but the tutorial code crashes with the error

import onnx_tf
[...many lines later...]
ModuleNotFoundError: No module named 'keras.src.engine'


which I have not been able to resolve yet because of dependency issues.

So my questions are:

1. Any ideas on how to fix the above mentioned Collab?
2. Has anyone successfully converted a YOLOX to .tflite and can point me in the right direction? There seems to be only so many ways to convert a model file and all of these seem to be a dead end of some sort.

Continue reading...