获取Tens部分对角线元素

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不允许axis，所以它在这里似乎没有用处。这是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慢的方式。但不确定浮点值的结果会是什么。另外，您可以尝试用tf.einsum替换tf.tensordot。关于第一个和第二个，它们看起来都很好，尽管如果通过这些操作进行反向传播，计算梯度的成本可能会有所不同。你知道吗