py_func只是用于CPU操作?

时间:2016-12-12 14:47:09

标签: tensorflow

根据文档,我可以使用tf.py_func来定义自己的操作。它将接收Numpy数组,并且必须返回Numpy数组。它没有明确地写在文档中,但我得出结论,它将被绑定到CPU设备?

如果我将其与可能在GPU上运行的其他操作相结合,TF会贪婪地向GPU移动尽可能多的计算并自动为我的tf.py_func操作系统在GPU和CPU之间传输内存吗? (就像Theano一样。)

是否有类似tf.py_func的类似内容来定义GPU操作?

1 个答案:

答案 0 :(得分:5)

py_func是一个奇怪的 - 它在与用于创建op的解释器相同的Python解释器中运行Python代码。因此,如果你周围的操作是GPU,那么将有GPU< - > CPU传输。此外,还会有一个副本在Python地址空间和TensorFlow地址空间之间移动数据(即memcpy here

如果你问如何在GPU上运行Python代码,那就是Numba。如果你在TensorFlow GPU中询问如何使用“类似功能”的东西,那就是Defun。一般来说,TensorFlow正朝着定义足够多的numpy特性的方向发展,这样你就不需要使用numpy,而是可以用原生的TF原语来实现你的函数。

在依赖TensorFlow引擎的同时,还有另一种潜在的Python-land方式。您可以创建一个实现模拟Python数字类型interface的Python包装器,但将实际工作委托给基础TensorFlow引擎。有点像numpy如何实现接口并将基础工作委托给BLAS库。

通过一些技巧,您可以将数据始终保留在GPU上。概述了概念实施概述here