最近项目,需要创建一个 tensorflow 的一个自定义操作,用来加速tensorflow的处理效果;下面对创建过程中,遇到的问题和资源进行简要记录,进行备忘:
OP 创建 参考链接:
https://www.tensorflow.org/guide/create_op (官方教程)
https://github.com/tensorflow/custom-op (官方模板,看完上面的教程,使用该模板就可以很方便得在docker 容器中进行尝试构建;较为推荐)
何时定义一个新的OP:
- 现有的operation 组合不出来需要的OP;
- 现有的operation 组合出来的OP 十分低效;
- 你想要融合一些操作来提高效率;
保持更新,转载请注明出处;更多内容请关注 cnblogs.com/xuyaowen;
安装测试docker(用于gpu环境docker测试):
docker pull tensorflow/tensorflow:custom-op-gpu-ubuntu16
docker run --runtime=nvidia --privileged -it -v ${PWD}:/working_dir -w /working_dir tensorflow/tensorflow:custom-op-gpu-ubuntu16
安装测试docker(用于cpu环境docker测试):
docker pull tensorflow/tensorflow:custom-op-ubuntu16
docker run -it -v ${PWD}:/working_dir -w /working_dir tensorflow/tensorflow:custom-op-ubuntu16
docker run -it tensorflow/tensorflow:custom-op-ubuntu16 /bin/bash
使用清华镜像临时下载pip安装包:
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple tensorflow