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

Increasing each element of a tensor by the predecessor in Tensorflow 2.0

我是 tensorflow 2.0 的新手,除了根据样板代码设计和训练一些人工神经网络外,没有做太多事情。我正在尝试将 新手 的练习解决到新的张量流中。我创建了一些代码,但它不起作用。下面是问题定义:


假设我们有形状为 (a, b, c) 的有理数张量 M 和标量 p ∈ (0, 1)(内存因子),让我们创建一个函数,它将 return 张量 N 的形状为 (a, b, c)。沿轴 c 移动的 N 个张量的每个元素应增加前任乘以 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 上提供最短的执行时间。创建的图应该允许有效地计算张量 M 和值 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.

我查看了文档,看到了 cumsumpolyeval 等函数,但我不确定它们是否符合我的需要。据我了解,我需要编写自己的客户函数并用 @tf.function 注释。我也不确定如何根据问题定义正确处理 3 维张量(添加前身应该发生在最后一个 ("c") 轴上)。

我在文档(此处:https://www.tensorflow.org/tutorials/customization/performance)中看到有一些方法可以测量生成的图形的大小。虽然,我不确定 "graph" 如何在张量 M 和值 p 上有效地计算 导数 。 ELI5 个答案表示赞赏,或者至少我可以阅读一些材料来更好地自学。

非常感谢!

我会给你几个不同的方法来实现它。我认为最明显的解决方案是使用 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 上可能会有很大差异。

这是一个仅提供一些信息的答案,以及修复代码的幼稚解决方案 --- 不是实际问题(请参阅下面的原因)。

首先,TypeError 是您早期尝试的张量中类型不兼容的问题。一些张量包含浮点数(双精度),一些包含整数。它会有助于显示 完整 错误消息:

TypeError: Input 'y' of 'Mul' Op has type int32 that does not match type float64 of argument 'x'.

这恰好放在了正确的轨道上(尽管堆栈跟踪的血淋淋的细节)。

这是使代码正常工作的简单修复(针对目标问题的警告):

import tensorflow as tf

@tf.function
def vectorize_predec(t, p):
    _p = tf.transpose(
        tf.convert_to_tensor(
            [p * t[...,idx] for idx in range(t.shape[-1] - 1)],
            dtype=tf.float64))
    _p = tf.concat([
        tf.zeroes((_p.shape[0], 1), dtype=tf.float64),
        _p
    ], axis=1)
    return t + _p

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]], dtype=tf.float64)

n = tf.constant([[0.0, 1.0, 2.5, 4.0, 5.5],
          [1.0, 3.5, 6.5, 9.5, 13.5],
          [1.0, 1.5, 1.5, -0.5, -0.5]], dtype=tf.float64)
print(f'Expected: {n}')

result = vectorize_predec(m, p)
print(f'Result: {result}')

tf.test.TestCase().assertAllEqual(n, result)

主要变化:

  • m 张量得到一个 dtype=tf.float64 来匹配原始的 double,因此类型错误消失了。
  • 功能基本完全重写。天真的想法是利用问题定义,它 而不是 说明 N 中的值是在更新之前还是之后计算的。这是更新前的版本,更简单。解决似乎是 "real" 的问题需要在函数上做更多的工作(参见其他答案,我可能会在这里做更多的工作)。

函数的工作原理:

  • 它将预期增量 p * x1p * x2 等计算到标准 Python 数组中。请注意它在最后一个维度的最后一个元素之前停止,因为我们将移动数组。
  • 它将数组转换为张量tf.convert_to_tensor,因此将数组添加到计算图中。转置是匹配原始张量形状所必需的(我们可以避免它)。
  • 它沿最后一个轴在每个维度的开头附加零。
  • 结果为原始张量与构造张量之和

值变为 x1 + 0.0 * p,然后是 x2 + x1 * p,等等。这说明了一些需要注意的功能和问题(类型,形状),但我承认它作弊并且没有解决实际问题问题。

此外,此代码在任何硬件上都不是高效的。这只是说明性的,需要 (1) 消除 Python 数组,(2) 消除转置,(3) 消除连接操作。希望有很好的培训:-)


补充说明:

  • 该问题要求求解形状为 (a, b, c) 的张量。您分享的代码适用于形状为 (a, b) 的张量,因此修复代码仍然无法解决问题。
  • 这道题需要有理数。不确定意图是什么 是的,这个答案把这个要求放在一边。
  • T = [x1, x2, x3, x4] 的形状实际上是 (4,),假设 xi 是标量。
  • 为什么tf.float64?默认情况下,我们得到 tf.float32,删除 double 将使代码正常工作。但是这个例子会失去类型重要的一点,所以选择一个明确的非默认类型(和更丑陋的代码)。