tf.function 比普通的 python 代码慢得多答案

【问题标题】：tf.function much slower than plain python codetf.function 比普通的 python 代码慢得多
【发布时间】：2021-02-07 21:43:50
【问题描述】：

为了提高我的项目的性能，我使用 tf.function 编写了一个函数来替换一个不使用 tf.function 的函数。结果是，在启用 GPU 时，纯 python 代码的运行速度比 tf.funtion 快得多（快 100 倍）。在 CPU 上运行时，TF 仍然较慢，但仅慢 10 倍。我错过了什么吗？

@tf.function
def test1(cond):
    xp = tf.constant(0)
    yp = tf.constant(0)
    stride = tf.constant(10)

    patches = tf.TensorArray(
        tf.int32, size=tf.cast((cond / stride + 1) * (cond / stride + 1), dtype=tf.int32), dynamic_size=False, clear_after_read=False)

    i = tf.constant(0)
    while tf.less_equal(yp, cond):
        while tf.less_equal(xp, cond):
            xp = tf.add(xp, stride)
            patches = patches.write(i, xp)
            i += 1
        xp = tf.constant(0)
        yp = tf.add(yp, stride)
    return patches.stack()

def test2(cond):
    xp = 0
    yp = 0
    stride = 10
    i = 0
    patches = []
    while yp <= cond:
        while xp <= cond:
            xp += stride
            patches.append(xp)
        xp = 0
        yp += stride
    return patches

当 cond 很大（例如 5000 或更大）时，这一点特别明显

更新：

我找到了this 和this。正如我所料，TensorArray 的性能似乎很差，而且，就我而言，解决方案是用其他张量计算替换 TensorArray 和循环（在这种情况下，我使用了 tf.image.extract_patches 等）。通过这种方式，它实现了比普通 python 代码快 3 倍的性能。

【问题讨论】：

标签： python tensorflow

【解决方案1】：

当您在启用 GPU 的情况下使用某些 tf 函数时，它会执行回调并将数据传输到 GPU。在某些情况下，这种开销是不值得的。在 CPU 上运行时，这种开销会减少，但仍然比纯 python 代码慢。

当您进行大量计算时，Tensorflow 会更快，而这正是 Tensorflow 的用途。即使是 numpy 也可能比纯 Python 代码更慢。

【讨论】：

原始代码有很多计算令我感到沮丧，因为问题似乎出在循环和 TensorArray 上。我知道 GPU 之间存在数据传输，但我不明白为什么它会这么慢，因为它甚至在 TF doc (TensorArray) 中被推荐。
你确定你有“很多计算”吗？如果您的程序只需几秒钟即可完成，请忘记使用 tensorflow。请记住，TensorFlow 用于在具有数千或数百万个参数的网络中进行 autograd。

【解决方案2】：

慢速部分（while 循环）仍在 python 中，像这样的简单函数非常快。每次从 python 切换到 tf 的线性开销肯定比你在这么小的函数上所能获得的任何东西都要大。对于更复杂的操作，这可能会非常不同。在这种情况下， tf 简直是矫枉过正。

【讨论】：

根据 TF doc.，while 也应该转换为 tf 因为条件是张量