在 Tensorflow 2.0 中通过前辈增加张量的每个元素

Question

我是tensorflow 2.0 的新手，除了从样板代码设计和训练一些人工神经网络之外，并没有做太多事情。我正在尝试为新人解决新 tensorflow 的练习。我创建了一些代码，但它不起作用。以下是问题定义：

---

假设我们有M形状为 的有理数张量(a, b, c)和标量p ? (0, 1)（记忆因子），让我们创建一个函数，该函数将返回N形状为 的张量(a, b, c)。N沿c轴移动的张量的每个元素都应增加前身乘以 的值p。

假设我们有张量：

T = [x1, x2, x3, x4]

在 的形状中(1, 1, 4)，我们想得到向量：

[x1, x2+x1·p, x3+(x2+x1·p)·p, x4+(x3+(x2+x1·p)·p)*p] 

解决方案应该在Tensorflow 2.0 中创建，并且应该专注于在 CPU 上提供最短的执行时间。创建的图形应该允许有效地计算 tensorM和 value 的导数p。

---

这是我创建到现在的代码：

import tensorflow as tf

@tf.function
def vectorize_predec(t, p):
    last_elem = 0
    result = []
    for el in t:
        result.append(el + (p * last_elem))
        last_elem = el + (p * last_elem)
    return result

p = tf.Variable(0.5, dtype='double')

m = tf.constant([[0, 1, 2, 3, 4],
          [1, 3, 5, 7, 10],
          [1, 1, 1, -1, 0]])

vectorize_predec(m, p)

但它抛出一个TypeError.

我查看了文档，我看到过cumsum和这样的函数polyeval，但我不确定它们是否符合我的需要。据我了解，我需要编写自己的客户函数，并用@tf.function. 我也不确定如何根据问题定义正确处理 3 维张量（添加前任应该发生在最后一个（“c”）轴上）。

我在文档（这里：https : //www.tensorflow.org/tutorials/customization/performance）中看到有一些方法可以测量生成的图形的大小。虽然，我不确定“图形”如何允许有效地计算tensor和 value 的导数。ELI5 答案表示赞赏，或者至少我可以阅读一些材料以更好地教育自己。Mp

非常感谢！

Answer 1

我会给你几个不同的方法来实现它。我认为最明显的解决方案是使用tf.scan：

import tensorflow as tf

def apply_momentum_scan(m, p, axis=0):
    # Put axis first
    axis = tf.convert_to_tensor(axis, dtype=tf.int32)
    perm = tf.concat([[axis], tf.range(axis), tf.range(axis + 1, tf.rank(m))], axis=0)
    m_t = tf.transpose(m, perm)
    # Do computation
    res_t = tf.scan(lambda a, x: a * p + x, m_t)
    # Undo transpose
    perm_t = tf.concat([tf.range(1, axis + 1), [0], tf.range(axis + 1, tf.rank(m))], axis=0)
    return tf.transpose(res_t, perm_t)

但是，如果您构建指数因子矩阵，您也可以将其实现为特定的矩阵乘积：

import tensorflow as tf

def apply_momentum_matmul(m, p, axis=0):
    # Put axis first and reshape
    m = tf.convert_to_tensor(m)
    p = tf.convert_to_tensor(p)
    axis = tf.convert_to_tensor(axis, dtype=tf.int32)
    perm = tf.concat([[axis], tf.range(axis), tf.range(axis + 1, tf.rank(m))], axis=0)
    m_t = tf.transpose(m, perm)
    shape_t = tf.shape(m_t)
    m_tr = tf.reshape(m_t, [shape_t[0], -1])
    # Build factors matrix
    r = tf.range(tf.shape(m_tr)[0])
    p_tr = tf.linalg.band_part(p ** tf.dtypes.cast(tf.expand_dims(r, 1) - r, p.dtype), -1, 0)
    # Do computation
    res_tr = p_tr @ m_tr
    # Reshape back and undo transpose
    res_t = tf.reshape(res_tr, shape_t)
    perm_t = tf.concat([tf.range(1, axis + 1), [0], tf.range(axis + 1, tf.rank(m))], axis=0)
    return tf.transpose(res_t, perm_t)

这也可以重写以避免第一次转置（在 TensorFlow 中很昂贵）tf.tensordot：

import tensorflow as tf

def apply_momentum_tensordot(m, p, axis=0):
    # Put axis first and reshape
    m = tf.convert_to_tensor(m)
    # Build factors matrix
    r = tf.range(tf.shape(m)[axis])
    p_mat = tf.linalg.band_part(p ** tf.dtypes.cast(tf.expand_dims(r, 1) - r, p.dtype), -1, 0)
    # Do computation
    res_t = tf.linalg.tensordot(m, p_mat, axes=[[axis], [1]])
    # Transpose
    last_dim = tf.rank(res_t) - 1
    perm_t = tf.concat([tf.range(axis), [last_dim], tf.range(axis, last_dim)], axis=0)
    return tf.transpose(res_t, perm_t)

这三个函数将以类似的方式使用：

import tensorflow as tf

p = tf.Variable(0.5, dtype=tf.float32)
m = tf.constant([[0, 1, 2, 3, 4],
                 [1, 3, 5, 7, 10],
                 [1, 1, 1, -1, 0]], tf.float32)
# apply_momentum is one of the functions above
print(apply_momentum(m, p, axis=0).numpy())
# [[ 0.    1.    2.    3.    4.  ]
#  [ 1.    3.5   6.    8.5  12.  ]
#  [ 1.5   2.75  4.    3.25  6.  ]]
print(apply_momentum(m, p, axis=1).numpy())
# [[ 0.      1.      2.5     4.25    6.125 ]
#  [ 1.      3.5     6.75   10.375  15.1875]
#  [ 1.      1.5     1.75   -0.125  -0.0625]]

使用矩阵乘积在渐近上更复杂，但它可以比扫描更快。这是一个小基准：

import tensorflow as tf
import numpy as np

# Make test data
tf.random.set_seed(0)
p = tf.constant(0.5, dtype=tf.float32)
m = tf.random.uniform([100, 30, 50], dtype=tf.float32)

# Axis 0
print(np.allclose(apply_momentum_scan(m, p, 0).numpy(), apply_momentum_matmul(m, p, 0).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 0).numpy(), apply_momentum_tensordot(m, p, 0).numpy()))
# True
%timeit apply_momentum_scan(m, p, 0)
# 11.5 ms ± 610 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit apply_momentum_matmul(m, p, 0)
# 1.36 ms ± 18.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit apply_momentum_tensordot(m, p, 0)
# 1.62 ms ± 7.39 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# Axis 1
print(np.allclose(apply_momentum_scan(m, p, 1).numpy(), apply_momentum_matmul(m, p, 1).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 1).numpy(), apply_momentum_tensordot(m, p, 1).numpy()))
# True
%timeit apply_momentum_scan(m, p, 1)
# 4.27 ms ± 60.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit apply_momentum_matmul(m, p, 1)
# 1.27 ms ± 36.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit apply_momentum_tensordot(m, p, 1)
# 1.2 ms ± 11.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# Axis 2
print(np.allclose(apply_momentum_scan(m, p, 2).numpy(), apply_momentum_matmul(m, p, 2).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 2).numpy(), apply_momentum_tensordot(m, p, 2).numpy()))
# True
%timeit apply_momentum_scan(m, p, 2)
# 6.29 ms ± 64.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit apply_momentum_matmul(m, p, 2)
# 1.41 ms ± 21.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit apply_momentum_tensordot(m, p, 2)
# 1.05 ms ± 26 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

所以，矩阵乘积似乎赢了。让我们看看这是否可以扩展：

import tensorflow as tf
import numpy as np

# Make test data
tf.random.set_seed(0)
p = tf.constant(0.5, dtype=tf.float32)
m = tf.random.uniform([1000, 300, 500], dtype=tf.float32)

# Axis 0
print(np.allclose(apply_momentum_scan(m, p, 0).numpy(), apply_momentum_matmul(m, p, 0).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 0).numpy(), apply_momentum_tensordot(m, p, 0).numpy()))
# True
%timeit apply_momentum_scan(m, p, 0)
# 784 ms ± 6.78 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_matmul(m, p, 0)
# 1.13 s ± 76.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_tensordot(m, p, 0)
# 1.3 s ± 27 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Axis 1
print(np.allclose(apply_momentum_scan(m, p, 1).numpy(), apply_momentum_matmul(m, p, 1).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 1).numpy(), apply_momentum_tensordot(m, p, 1).numpy()))
# True
%timeit apply_momentum_scan(m, p, 1)
# 852 ms ± 12.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_matmul(m, p, 1)
# 659 ms ± 10.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_tensordot(m, p, 1)
# 741 ms ± 19.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Axis 2
print(np.allclose(apply_momentum_scan(m, p, 2).numpy(), apply_momentum_matmul(m, p, 2).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 2).numpy(), apply_momentum_tensordot(m, p, 2).numpy()))
# True
%timeit apply_momentum_scan(m, p, 2)
# 1.06 s ± 16.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_matmul(m, p, 2)
# 924 ms ± 17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_tensordot(m, p, 2)
# 483 ms ± 10.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

好吧，现在它不再那么清楚了。扫描仍然不是超快，但矩阵产品有时更慢。可以想象，如果使用更大的张量，矩阵乘积的复杂性将主导时序。

因此，如果您想要最快的解决方案并且知道您的张量不会变得很大，请使用矩阵乘积实现之一。如果您对速度还算满意，但想确保不会耗尽内存（矩阵解决方案也需要更多时间）并且时间是可预测的，您可以使用扫描解决方案。

注意：以上基准测试是在 CPU 上进行的，结果在 GPU 上可能会有很大差异。