Skip to content

bo-dun/tensorflow-tensorrt

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

19 Commits
 
 
 
 
 
 
 
 

Repository files navigation

Tensorflow to TensorRT Model Converter

This conversion implementation uses a JSON intermediate to port neural network information. Any TensorFlow model can be converted to TensorRT format, optimized, and used for highly efficient inferencing. The current converter from JSON to TensorRT includes customized code for loading ImageNet data for inferencing.

Files

Tensorflow-JSON.py takes a TensorFlow Session and a few other related objects and serializes them to JSON format. In the code I included two variants on the serialization method, one that aggregates the entire network into one JSON object and one that serializes each layer individually. When using, make sure that the dictionary is formatted correctly according to my comments in the file.

JSON-TensorRT.cpp deserializes a directory of JSON layer files into a TensorRT model. This file also includes customized code to load and synchronously evaluate ImageNet data for direct inferencing. If the input data you're loading has multiple channels and/or batches, it must be in NCHW format, which is what TensorRT works with. TensorFlow, however, uses NHWC format. When debugging, keep in mind that convolutional weights are in KCRS format in TensorRT and RSCK format in TensorFlow.

Author

Feel free to reach out with any questions, comments, suggestions, bug reports, etc.

License

This project is licensed under the MIT License - see the LICENSE file for more details.

Releases

No releases published

Packages

No packages published