Google Cloud ML引擎中的TensorFlow新OP

Tensorflow new op in google cloud ml engine

本文关键字：TensorFlow OP Cloud ML 引擎 Google      更新时间：2023-10-16

我创建了一个新的TensorFlow OP，我已经在本地进行了编译并进行了测试。

我现在想与Google Cloud ML引擎一起使用，该发动机需要在云中的每台机器中编译OP。

我设法使用setup.py：

将C 文件和编译器脚本包括在分发中 
from setuptools import setup, find_package
setup(
    name='trainer',
    version='0.1',
    packages=find_packages(),
    package_data={'CPP': ['*.cc', '*.h', 'compile.sh']},
    description = 'Package discription'
)

现在我必须运行包含：

的compile.sh 
TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')
g++ -std=c++11 -D_GLIBCXX_USE_CXX11_ABI=0 -shared target.cc dependency_1.cc -o taget.so -fPIC -I$TF_INC -I$TF_INC/external/nsync/public -O2

可以使用subprocess.Popen()从Python代码运行脚本，因此这不是问题。

问题是我不知道目录树在Google Cloud中是如何构造的，因此我不知道该脚本在哪里运行，或者以后如何访问其输出以利用新OP。

最简单的方法是在本地构建OP，上传到您的GCS存储桶，将其复制到VM（容器），然后通过TF.Load_op_library使用它。您可以使用预装的GSUTIL CLI或GCS PYTHON客户端执行副本。

回到原始问题：当我们启动用户作业时，我们首先将用户代码包作为根，因此在Python 2.7的情况下，它位于/root/.local/lib/python2.7/site-packages/your_package_name

这是我的解决方案。我使用setup的Extension模块在云中部署软件包时动态构建OP：

import tensorflow as tf
from setuptools import setup, find_packages, Extension

TF_INC = tf.sysconfig.get_include()
TF_LIB = tf.sysconfig.get_lib()
module1 = Extension('op_name',
                    define_macros=[('_GLIBCXX_USE_CXX11_ABI', 0)],
                    include_dirs=[TF_INC, TF_INC + '/external/nsync/public', 'Op/cpp'],
                    sources=['Op/cpp/op_name.cc',
                             'Op/cpp/source_2.cc'],
                    extra_compile_args=['-O2', '-std=c++11', '-fPIC', '-shared'],
                    library_dirs=[TF_LIB],
                    libraries=['tensorflow_framework'],
                    extra_link_args=[],
                    language='c++')
setup(
    name='trainer',
    version='0.1',
    packages=find_packages(),
    package_data={'Op': ['cpp/*.cc', 'cpp/*.h']},
    ext_modules=[module1],
)

一些笔记：

  • 如果您有标题文件，则必须在include_dirs中包含这些文件的目录。在这种情况下，我将标题文件与源文件相同的目录（Op/cpp）。
  • 但是，这并不意味着.h文件已包装。为此必须使用package_data={'Op': ['cpp/*.cc', 'cpp/*.h']},，以便 .h文件包含在清单中。.cc文件无论如何都应包括在内，因为它们是来源，我只是在这里，因为。
  • 云中使用的编译器是使用c1plus的gcc。TensorFlow官方文档使用g++。我不知道这对性能有什么影响...
相关文章：