google 云机器学习引擎中的 Tensorflow 新操作

Question

我创建了一个新的 Tensorflow 操作，我已经在本地编译和测试它并且可以运行。

我现在想将它与 Google Cloud ML 引擎一起使用，这需要在云中的每台机器上编译操作。

我已经成功地使用 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

脚本可以运行来自 python 代码，使用 subprocess.Popen() 所以这不是问题。

问题是我不知道目录树在 google 云中的结构，所以我不知道运行这个脚本的位置，或者如何访问它的输出稍后使用新的操作。

Answer 1

最简单的方法是在本地构建 op，上传到您的 GCS 存储桶，将其复制到 VM（容器），然后通过 tf.load_op_library 使用它。您可以使用预安装的 gsutil cli 或 GCS python 客户端来执行复制。

回到最初的问题：当我们开始一个用户作业时，我们首先以 root 身份安装用户代码包，所以在 python 2.7 的情况下，它位于 /root/.local/lib/python2.7/site-packages/YOUR_PACKAGE_NAME

Answer 2

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

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 文件无论如何都应该包括在内，因为它们是源代码，我把它们放在这里是因为。
云中使用的编译器是gcc，它使用c1plus。 Tensorflow官方文档使用g++。我不知道这对性能有什么影响...

google 云机器学习引擎中的 Tensorflow 新操作

Tensorflow new op in google cloud ml engine

c++

python

tensorflow

google-cloud-ml