This sample demonstrates how to first train a model using TensorFlow and Keras, freeze the model and write it to a protobuf file, convert it to UFF, and finally run inference using TensorRT5.
-
Make sure you have the python dependencies installed.
- For python2, run
python2 -m pip install -r requirements.txt
from the top-level of this sample. - For python3, run
python3 -m pip install -r requirements.txt
from the top-level of this sample.
- For python2, run
-
Make sure you have the UFF toolkit as well as
graphsurgeon
installed. if not, download it from here -
Train the model and write out the frozen graph:
$ mkdir models $ python train.py
-
Convert the .pb file to .uff, using the convert-to-uff utility:
$ convert-to-uff ./models/model.pb -o ./models/
The converter will display information about the input and output nodes, which you can use to the register inputs and outputs with the parser. In this case, we already know the details of the input and output nodes and have included them in the sample.
-
Create a TensorRT inference engine from the uff file
$ python uff2plan.py
Create a TensorRT inference engine from the pb file:
$ python pb2plan.py
$ python inference.py