python - 使用 tensorflow 一次切片多个切片

Question

我正在为 tensorflow RNN 准备输入张量。
目前我正在做以下事情

rnn_format = list()
for each in range(batch_size):
    rnn_format.append(tf.slice(input2Dpadded,[each,0],[max_steps,10]))
lstm_input = tf.stack(rnn_format)

是否可以在没有循环的情况下使用一些 tensorflow 函数立即执行此操作？

Answer 1

正如 Peter Hawkins 所建议的，您可以使用 gather_nd 和适当的索引来到达那里。

您可以在调用 gather_nd 之前简单地对内部维度进行统一裁剪。

例子:

import tensorflow as tf
import numpy as np

sess = tf.InteractiveSession()

# integer image simply because it is more readable to me
im0 = np.random.randint(10, size=(20,20))
im = tf.constant(im0)

max_steps = 3
batch_size = 10

# create the appropriate indices here
indices = (np.arange(max_steps) +
    np.arange(batch_size)[:,np.newaxis])[...,np.newaxis]
# crop then call gather_nd
res = tf.gather_nd(im[:,:10], indices).eval()

# check that the resulting tensors are equal to what you had previously
for each in range(batch_size):
  assert(np.all(tf.slice(im, [each,0],[max_steps,10]).eval() == res[each]))

编辑

如果您的切片索引在张量中，您只需在创建 indices 时将 numpy 的操作替换为 tensorflow 的操作:

# indices stored in a 1D array
my_indices = tf.constant([1, 8, 3, 0, 0])
indices = (np.arange(max_steps) +
    my_indices[:,tf.newaxis])[...,tf.newaxis]

进一步说明:

• indices 是利用 broadcasting 创建的添加期间:数组实际上是平铺的，因此它们的尺寸匹配。 numpy 和 tensorflow 以类似的方式支持广播。
• 省略号 ... 是标准的一部分 numpy slicing notation ，它基本上填充了其他切片索引留下的所有剩余维度。所以 [..., newaxis] 基本上等同于 expand_dims(·, -1)。