谷歌云ml引擎中的Tensorflow新操作

Posted 2023-02-23

【中文标题】谷歌云ml引擎中的Tensorflow新操作

【英文标题】：Tensorflow new op in google cloud ml engine

【发布时间】：2017-12-19 11:12:25

【问题描述】：

我创建了一个新的 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

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

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

【参考方案1】：

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

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

【讨论】：

如果我错了，请纠正我，但如果我在本地编译操作，它将无法在 gcloud 机器上运行。对？我使用的是 macOS，而 gcloud 使用的是 linux。

是的，我们使用的是 Linux，在 mac 上构建的操作可能无法正常工作。如果您仍然想尝试我建议的第一个，您可以在 GCE 上的 Linux VM 中构建操作。

好的，我试试。你以前做过吗？ （在 GCE 虚拟机中构建，然后使用 op）。

不，我之前没有这样做过，但是VM与本地机器没有区别，所以它应该可以工作。

【参考方案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++。我不知道这对性能有什么影响...