Skip to content
This repository has been archived by the owner on Nov 29, 2022. It is now read-only.

artemmavrin/TF-Custom-CXX-Op

Repository files navigation

TensorFlow Custom C++ Op

Build Last Commit License

This Repo Is Under Construction

This is an example of how to create a custom TensorFlow op in C++ and use it in Python.

The op implemented here is the logit function, logit(x)=log(x/(1-x)) component-wise for x with entries between 0 and 1 (inclusive). This is the inverse of the sigmoid function, and can be used to turn a probability into a score (the log odds) between -inf and inf. The logit function is currently not implemented as a native TensorFlow op (a NumPy ndarray ufunc version exists in SciPy as scipy.special.logit).

Usage

Clone this repo and run make:

git clone https://github.com/artemmavrin/TF-Custom-CXX-Op.git
cd TF-Custom-CXX-Op
# Optional but recommended: create and activate a new Python virtual environment
make

Running make will download the required Python dependencies, install the tf_custom_cxx_op Python package, and compile the C++ op. It will also run the unit tests. These steps can be done separately using make install (to install dependencies and the Python package), make build (to compile the C++ op), and make test (to run the unit tests).

Using the logit TensorFlow op in Python is then easy. For example, in a Python REPL:

>>> from tf_custom_cxx_op import logit
>>> logit(0.5)  # Accepts a tensor of any shape
<tf.Tensor: shape=(), dtype=float32, numpy=0.0>
>>> logit([0.0, 0.25, 0.5, 0.75, 1.0])
<tf.Tensor: shape=(5,), dtype=float32, numpy=
array([      -inf, -1.0986123,  0.       ,  1.0986123,        inf],
      dtype=float32)>

TODO

  • Add documentation.
  • Test on macOS.
  • Test on Linux.
  • Test on Windows.
  • Create pip-installable package.
  • Register gradient directly in C++.
  • Make a GPU op.

Acknowledgements

Good reading about creating TensorFlow C++ ops: