Jupyter Notebook to convert TF/Keras Model to TenosrRT with ONNX

tf2trt_with_onnx this Github repo has Jupyter notebook documents how to convert a Tensorflow or Keras model to TenosrRT using ONNX. ONNX is the official way to convert TF/Keras model to TRT.

This is has been tested successfully on Facenet Keras model
If anyone has a question please ask!

1 Like

Hi, plz could you help me, I have clone the git repository into my tensorflow nvidia container, TF2.2-nv. So, I have made some changes: instead of import keras, i used from tensorflow import keras. That is step 1.
Then, in step 2, I got, module ‘tensorflow.keras.backend’ has no attribute ‘get_session’.
So, I would really appreciate if you recommend to downgrade my TF. or to use another container in order to be able to execute this Jupyter notebook.

In TF 2 get_session is removed, there are some modifications needed, I have added keras_to_pb_tf2.py
I have test it successfully to generate pb file in TF2.0. In jupyter notebook just change the from keras_to_pb_tf2 import keras_to_pb' instead of from keras_to_pb import keras_to_pb’

hope it works…


Thanks a lot, it worked for me (well, I ran two first steps in other laptop, with those files y could continue with step 3 in my new container, cause I had some trouble and I reconstruct my docker container). Finally, I have my TRT engine. I hope it works on inference =). Ill let you know.


Now i can make the inferences with my facenet with tensorRT, but it seems i am not using the gpu =(. It is not being faster.

Step 2 convert pb to ONNX must be done on same machine that you will do inferencing because it will optimize the model specificaly to the GPU on the machine. Did you do it on same machine?


oh I did it in my laptop, because in the same machine tensorflow just gave me errors. (TF2.2 on nvidia dgx-a100, new toy ja). Ill keep trying. Thanks.

1 Like

Now i am repeating the process… tks, There are no error in uploading the files (also i had to uploaded again, weird).

1 Like

I’m getting this error:- ModuleNotFoundError: No module named ‘tensorrt’

while running cmd:- from onnx_to_trt import create_engine

onnx_to_trt goes to engine.py file and their we import tensorrt, use dont have that file.