是否有更简单的方法来获取张量的切片，如以下示例所示？

Question

我想像下面的numpy中那样对张量进行切片。我该怎么办？

# numpy array
a = np.reshape(np.arange(60), (3,2,2,5))
idx = np.array([0, 1, 0])
N = np.shape(a)[0]
mask = a[np.arange(N),:,:,idx]


# I have tried several solutions, but only the following success.
# tensors
import tensorflow as tf
import numpy as np


a = tf.cast(tf.constant(np.reshape(np.arange(60), (3,2,2,5))), tf.int32)
idx2 = tf.constant([0, 1, 0])

fn = lambda i: a[i][:,:,idx2[i]]
idx = tf.range(tf.shape(a)[0])
masks = tf.map_fn(fn, idx)
with tf.Session() as sess:
    print(sess.run(a))
    print(sess.run(tf.shape(masks)))
    print(sess.run(masks))

有没有更简单的方法来实现这一目标？

我可以使用功能tf.gather或tf.gather_nd来实现吗？ 非常感谢！

Answer 1

1。另一种方法

我不确定这是否是最好的方法，但是速度更快。您可以使用val cols_needed = "a.col1,a.col2,a.col3,a.col4,b.col1,b.col2".split(",") for(c <- cols_needed) { if(Try(df(c)).isFailure) { df.withColumn(c, lit(null)) } else { df.withColumn(`c`,df(c)) } } 代替tf.boolean_mask。

tf.map_fn

2。性能比较

在我的8 GB GPU内存上，向量化方法进行1000次迭代需要import tensorflow as tf import numpy as np a = tf.cast(tf.constant(np.reshape(np.arange(60), (3,2,2,5))), tf.int32) idx2 = tf.constant([0, 1, 0]) fn = lambda i: a[i,:,:][:,:,idx2[i]] idx = tf.range(tf.shape(a)[0]) masks = tf.map_fn(fn, idx) # new method idx = tf.one_hot(idx2,depth=a.shape[-1]) masks2 = tf.boolean_mask(tf.transpose(a,[0,3,1,2]), idx) with tf.Session() as sess: print('tf.map_fn version:\n',sess.run(masks)) print('tf.boolean_mask version:\n',sess.run(masks2)) # print tf.map_fn version: [[[ 0 5] [10 15]] [[21 26] [31 36]] [[40 45] [50 55]]] tf.boolean_mask version: [[[ 0 5] [10 15]] [[21 26] [31 36]] [[40 45] [50 55]]] ，而0.07s方法1000次迭代需要tf.map_fn。向量化方法将比0.85s快得多。

tf.map_fn()

我相信，随着import datetime ... with tf.Session() as sess: start = datetime.datetime.now() for _ in range(1000): sess.run(masks) end = datetime.datetime.now() print('tf.map_fn version cost time(seconds) : %.2f' % ((end - start).total_seconds())) start = datetime.datetime.now() for _ in range(1000): sess.run(masks2) end = datetime.datetime.now() print('tf.boolean_mask version cost time(seconds) : %.2f' % ((end - start).total_seconds())) # print tf.map_fn version cost time(seconds) : 0.85 tf.boolean_mask version cost time(seconds) : 0.07 形状的增加，性能差异会变得更加明显。

Answer 2

另一种方法使用tf.gather_nd：

import tensorflow as tf
import numpy as np


a = tf.cast(tf.constant(np.reshape(np.arange(60), (3,2,2,5))), tf.int32)
idx = tf.range(tf.shape(a)[0])
idx2 = tf.constant([0,1,0])
indices = tf.stack([idx, idx2], axis=1)
a = tf.transpose(a, [0,3,1,2])
masks = tf.gather_nd(a, indices)

with tf.Session() as sess:
    print(sess.run(a))
    print(sess.run(tf.shape(masks)))
    print(sess.run(masks))