Theano erweiterte Indizierung für Tensor, geteilter Index

Question

Ich habe einen Tensor probs mit probs.shape = (max_time, num_batches, num_labels).

Und ich habe einen Tensor targets mit targets.shape = (max_seq_len, num_batches), wo die Werte sind Label Indizes, d.h. für die dritte Dimension in probs.

Jetzt möchte ich einen Tensor probs_y mit probs.shape = (max_time, num_batches, max_seq_len) wo die dritte Dimension ist der Index in targets. Grundsätzlich

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

für alle 0 <= i < num_batches.

Wie kann ich das erreichen?

Ein ähnliches Problem mit der Lösung wurde here veröffentlicht.

Die Lösung gibt, wenn ich richtig verstehe, wäre:

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

Aber das scheint nicht zu funktionieren. Ich bekomme: IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices.

Auch ist die Erstellung der zeitlichen T.arange ein bisschen teuer? Esp, wenn ich versuche, Abhilfe zu schaffen, indem ich es wirklich ein volles dichtes Integer-Array mache. Es sollte einen besseren Weg geben.

Vielleicht theano.map? Aber soweit ich das verstehe, wird der Code nicht parallelisiert, also ist dies auch keine Lösung.

Answer 1

Dies funktioniert für mich:

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() 

Above verwendet eine transponierte Version der Indizes. Ihr genauer Vorschlag funktioniert auch

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() 

Also vielleicht ist dein Problem woanders.

Was die Geschwindigkeit angeht, bin ich nicht sicher, was schneller sein könnte. map, die eine Spezialisierung von scan ist, ist fast sicher nicht. Ich weiß nicht, in welchem ​​Ausmaß die arange tatsächlich gebaut und nicht einfach über iteriert wurde.