Question

我有一个张量元素，其形状为（？，a，a，b）。我想将其转换为形状张量（？，a，b），其中：

output[ i , j , k ] = input[ i , j , j , k ].

这在numpy中很简单，因为我可以通过遍历i，j，k来分配元素。但是，所有操作必须按其在Tensorflow中的张量保持不变，以评估成本函数和训练模型。

我已经看过tf.diag_part()，但是据我了解，这不能在特定轴上指定，必须在整个张量上完成。

Answer 1

由于您说tf.diag_part不允许使用轴，因此在这里似乎没有用。这是tf.gather_nd的一种可能的解决方案：

import tensorflow as tf
import numpy as np

# Input data
inp = tf.placeholder(tf.int32, [None, None, None, None])
# Read dimensions
s = tf.shape(inp)
a, b, c = s[0], s[1], s[3]
# Make indices for gathering
ii, jj, kk = tf.meshgrid(tf.range(a), tf.range(b), tf.range(c), indexing='ij')
idx = tf.stack([ii, jj, jj, kk], axis=-1)
# Gather result
out = tf.gather_nd(inp, idx)

# Test
with tf.Session() as sess:
    inp_val = np.arange(36).reshape(2, 3, 3, 2)
    print(inp_val)
    # [[[[ 0  1]
    #    [ 2  3]
    #    [ 4  5]]
    # 
    #   [[ 6  7]
    #    [ 8  9]
    #    [10 11]]
    # 
    #   [[12 13]
    #    [14 15]
    #    [16 17]]]
    # 
    # 
    #  [[[18 19]
    #    [20 21]
    #    [22 23]]
    # 
    #   [[24 25]
    #    [26 27]
    #    [28 29]]
    # 
    #   [[30 31]
    #    [32 33]
    #    [34 35]]]]
    print(sess.run(out, feed_dict={inp: inp_val}))
    # [[[ 0  1]
    #   [ 8  9]
    #   [16 17]]
    # 
    #  [[18 19]
    #   [26 27]
    #   [34 35]]]

这里有几个替代版本。一种使用张量代数。

inp = tf.placeholder(tf.int32, [None, None, None, None])
b = tf.shape(inp)[1]
eye = tf.eye(b, dtype=inp.dtype)
inp_masked = inp * tf.expand_dims(eye, 2)
out = tf.tensordot(inp_masked, tf.ones(b, inp.dtype), [[2], [0]])

还有一个使用布尔掩码：

inp = tf.placeholder(tf.int32, [None, None, None, None])
s = tf.shape(inp)
a, b, c = s[0], s[1], s[3]
mask = tf.eye(b, dtype=tf.bool)
inp_mask = tf.boolean_mask(inp, tf.tile(tf.expand_dims(mask, 0), [a, 1, 1]))
out = tf.reshape(inp_mask, [a, b, c])

编辑：我对这三种方法进行了一些时间测量：

import tensorflow as tf
import numpy as np

def f1(inp):
    s = tf.shape(inp)
    a, b, c = s[0], s[1], s[3]
    ii, jj, kk = tf.meshgrid(tf.range(a), tf.range(b), tf.range(c), indexing='ij')
    idx = tf.stack([ii, jj, jj, kk], axis=-1)
    return tf.gather_nd(inp, idx)

def f2(inp):
    b = tf.shape(inp)[1]
    eye = tf.eye(b, dtype=inp.dtype)
    inp_masked = inp * tf.expand_dims(eye, 2)
    return tf.tensordot(inp_masked, tf.ones(b, inp.dtype), [[2], [0]])

def f3(inp):
    s = tf.shape(inp)
    a, b, c = s[0], s[1], s[3]
    mask = tf.eye(b, dtype=tf.bool)
    inp_mask = tf.boolean_mask(inp, tf.tile(tf.expand_dims(mask, 0), [a, 1, 1]))
    return tf.reshape(inp_mask, [a, b, c])

with tf.Graph().as_default():
    inp = tf.constant(np.arange(100 * 300 * 300 * 10).reshape(100, 300, 300, 10))
    out1 = f1(inp)
    out2 = f2(inp)
    out3 = f3(inp)
    with tf.Session() as sess:
        v1, v2, v3 = sess.run((out1, out2, out3))
        print(np.all(v1 == v2) and np.all(v1 == v3))
        # True
        %timeit sess.run(out1)
        # CPU: 1 ms ± 138 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
        # GPU: 1.04 ms ± 93.7 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
        %timeit sess.run(out2)
        # CPU: 1.17 ms ± 150 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
        # GPU: 734 ms ± 17.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
        %timeit sess.run(out3)
        # CPU: 1.11 ms ± 172 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
        # GPU: 1.41 ms ± 1.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

似乎所有这三个在CPU上都是相似的，但是第二个由于某种原因在我的GPU中会变慢。不知道会有float值的结果是什么。此外，例如，您可以尝试将tf.tensordot替换为tf.einsum。关于第一个和第二个，它们似乎都不错，尽管如果通过这些操作进行反向传播，则计算梯度的成本可能会有所不同。

获取张量的仅一部分的对角线元素

1 个答案: