python - Tensorflow:如何在维数不变的情况下对张量进行切片？

Question

例如，如果我们有:

a = tf.constant(np.eye(5))
a
<tf.Tensor 'Const:0' shape=(5, 5) dtype=float64>
a[0,:]
<tf.Tensor 'strided_slice:0' shape=(5,) dtype=float64>

张量a的切片会将原来的维度数2减少到1

我怎样才能直接获得排名不改变的切片，例如:？

a[0,:]
<tf.Tensor 'strided_slice:0' shape=(1,5) dtype=float64>

(tf.expand_dims(a[0,:], axis=0) 可以工作，但是有更直接和简单的方法吗？)

Answer 1

至少有两种直接的方法，与 NumPy 中提供的方法非常相似 ( related question )。

1. 获取该轴上大小为 1 的范围:a[x:x+1]
2. 添加带有None的轴:a[None, x]

a[0:1]

<tf.Tensor 'strided_slice_1:0' shape=(1, 5) dtype=float64>

一些实际的张量运行显示了预期的结果。

with tf.Session() as sess:
    sess.run(a[0])
    sess.run(a[0:1])
    sess.run(a[None, 0])

array([1., 0., 0., 0., 0.])
array([[1., 0., 0., 0., 0.]])
array([[1., 0., 0., 0., 0.]])