用于张量、共享索引的 Theano 高级索引

Question

我有一个张量 probs 和 probs.shape = (max_time, num_batches, num_labels)。

我有一个张量targets 和targets.shape = (max_seq_len, num_batches)，其中的值是标签索引，即probs 中的第三维。

现在我想得到一个张量probs_y 和probs.shape = (max_time, num_batches, max_seq_len)，其中第三维是targets 中的索引。基本上

probs_y[:,i,:] = probs[:,i,targets[:,i]]

所有0 <= i < num_batches。

我怎样才能做到这一点？

here 发布了类似的解决方案。

如果我理解正确的话，那里的解决方案是：

probs_y = probs[:,T.arange(targets.shape[1])[None,:],targets]

但这似乎不起作用。我得到： IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices.

另外，时间T.arange 的创建不是有点昂贵吗？特别是当我尝试通过真正使其成为一个完整的密集整数数组来解决问题时。应该有更好的办法。

也许theano.map？但据我了解，这不会并行化代码，所以这也不是解决方案。

Answer 1

这对我有用：

import theano
import theano.tensor as T

max_time, num_batches, num_labels = 3, 4, 6
max_seq_len = 5

probs_ = np.arange(max_time * num_batches * num_labels).reshape(
    max_time, num_batches, num_labels)

targets_ = np.arange(num_batches * max_seq_len).reshape(max_seq_len, 
    num_batches) % (num_batches - 1)  # mix stuff up

probs, targets = map(theano.shared, (probs_, targets_))

print probs_
print targets_

probs_y = probs[:, T.arange(targets.shape[1])[:, np.newaxis], targets.T]

print probs_y.eval()

以上使用了您的索引的转置版本。你的确切提议也有效

probs_y2 = probs[:, T.arange(targets.shape[1])[np.newaxis, :], targets]

print probs_y2.eval()
print (probs_y2.dimshuffle(0, 2, 1) - probs_y).eval()

所以也许你的问题出在其他地方。

至于速度，我不知道还有什么比这更快。 map，这是 scan 的一个专业化，几乎可以肯定不是。我不知道arange 实际构建到什么程度，而不是简单地迭代。