python - 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吗？有点贵？ Esp 当我尝试通过真正使它成为一个完整的密集整数数组来解决问题时。应该有更好的方法。

也许 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 实际构建到什么程度，而不是简单地迭代。