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
BUILD
file in thesrc/
directory:- Change the
-i
and-o
options 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-NAME
is the full-path name of one of the test cases or suites.Run
bazel test //test:all
to 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:compiler
by 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 -h
in 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