Thanks to visit codestin.com
Credit goes to github.com

Skip to content

aniskhan25/custom-op

 
 

Repository files navigation

TensorFlow Custom Op

Build Example zero_out Op (CPU only)

Clone the repository.

git clone https://github.com/tensorflow/custom-op.git
cd custom-op

Build PIP Package

You can build the pip package with make.

  make -f Makefile.puhti zero_out_pip_pkg

Install and Test PIP Package

Once the pip package has been built, you can install it with,

pip3 install artifacts/*.whl

Then test out the pip package

cd ..
python3 -c "import tensorflow as tf;import tensorflow_zero_out;print(tensorflow_zero_out.zero_out([[1,2], [3,4]]))"

And you should see the op zeroed out all input elements except the first one:

[[1 0]
 [0 0]]

Build Example time_two Op (GPU only on Puhti)

Clone the repository.

git clone https://github.com/tensorflow/custom-op.git
cd custom-op

Load module

Load module,

ml tensorflow/2.12

Set env,

export SINGULARITY_BIND=$PWD

Run shell within the container

singularity shell $SING_IMAGE

Build the op with make,

make -f Makefile.puhti time_two_op

Test the op,

make -f Makefile.puhti time_two_test

Build Example time_two Op (GPU only on Lumi)

Clone the repository.

git clone https://github.com/tensorflow/custom-op.git
cd custom-op

Load module

Load module,

module use /appl/local/csc/modulefiles/
ml tensorflow/2.12

Set env,

export SINGULARITY_BIND=$PWD

Run shell within the container

singularity shell $SING_IMAGE

Build the op with make,

make -f Makefile.lumi time_two_op

Test the op,

make -f Makefile.lumi time_two_test

About

Guide for building custom op for TensorFlow

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Smarty 46.0%
  • Starlark 26.3%
  • Python 19.2%
  • Shell 4.5%
  • C++ 4.0%