python - Theano张量共享索引高级索引

我有一个带有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？但是据我了解，这不会并行化代码，因此这也不是解决方案。

最佳答案

这对我有用:

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