Compiles a pre-trained TensorFlow model to a C++ library and provides an inference function that uses the learned weights of the neural network model as constants.
The currently supported operations are Add, Sub, Mul, Div, Matmul, their respective broadcasting versions and some activation functions like ReLU, ReLU6 and LeakyReLU.
The last release is v0.3 and this repository is now archived.
src/: The source code directory that contains the TF compiler scripts.
src/templates/: The directory containing the Jinja C++ templates for the library that is to be generated.
test/: Contains tests for every supported operation for small TF models.
third_party/: The third-party dependency files. (Eigen, GTest, CUDA)
-
Make sure you have a ProtoBuf model. Check the freeze_graph utility in TF to see how you can freeze your model.
(Some graph-freezing utility is also available in the test scripts under
test/util.py) -
Familiarize yourself with the usage of the compiler by invoking
python src/main.py -h. -
Change the
BUILDfile in thesrc/directory:- Change the
-iand-ooptions by writing the input and output node names in the respective places.
- Change the
-
To generate the library files (.h and .cc) run:
bazel build //src:compiler
-
To test the library run:
bazel test //test/<TEST-NAME>
where
TEST-NAMEis the full-path name of one of the test cases or suites.Run
bazel test //test:allto perform all tests. -
To see the dependency graph for the Brook library:
bazel query --noimplicit_deps 'deps(//src:compiler)' --output graph > simplified_graph.in dot -Tpng < simplified_graph.in > graph.png
where you can substitute
//src:compilerby the name of the rule whose graph you would like to generate.
- When the compiler code is changed, the test models should be re-generated and tests should be re-run. For that purpose, you can use the script utility
generate_and_test_all_models.sh -hin thetest/directory. The help message should be self-explanatory. (From root directory runtest/generate_and_test_all_models.sh -h) - The
src/lib/directory doesn't exist initially. Therefore, in rare occasions, Bazel may complain the the build has failed, when it has, indeed, generated all the necessary files. This is a known issue with Bazel, documented in #3114
- Extending the range of supported operations
- Extending matrix algebra to Tensor algebra
- Compression:
- Quantization
- Low precision GEMM