インデックス

Question

をindex_select使用して別ずつPyTorchテンソルIは3×3 PyTorch LongTensorを持っている：

A = 
    [0, 0, 0] 
    [1, 2, 2] 
    [1, 2, 3] 

私が欲しい私たちには、インデックスに4×2 FloatTensorこのような：

B = 
    [0.4, 0.5] 
    [1.2, 1.4] 
    [0.8, 1.9] 
    [2.4, 2.9] 

私の意図した出力は、以下の2×3×3 FloatTensorある：行列、すなわち

C[0,:,:] = 
    [0.4, 0.4, 0.4] 
    [1.2, 0.8, 0.8] 
    [1.2, 0.8, 2.4] 

C[1,:,:] = 
    [0.5, 0.5, 0.5] 
    [1.4, 1.9, 1.9] 
    [1.4, 1.9, 2.9] 

Aは索引付けおよびブロードキャストマトリックスBです。 Aは、Bのインデックスの行列です。したがって、この操作は本質的にインデックス作成操作です。

torch.index_select()機能を使用してこれを行うにはどうすればよいですか？ソリューションにディメンションを追加または置換する必要がある場合は問題ありません。

Answer 1

index_select()を使用すると、インデックス値がテンソルではなくベクトルになる必要があります。しかし、それが正しくフォーマットされている限り、関数はあなたのために放送を処理します。最後に行わなければならないことは、出力を再構成することです。私は放送のためだと信じています。

正常にこの操作を行いますワンライナーは

torch.index_select(B, 0, A.view(-1)).view(3,-1,2).permute(2,0,1) 

A.view(-1)指数行列をベクトル化しています。

__.view(3,-1,2)は、インデックス行列の形状に戻りますが、サイズ2の新しい余分なディメンション（N×2行列のインデックスを作成しています）を反映しています。

最後に__.permute(2,0,1)は、各出力がBの各ディメンションを（各列ではなく）別々のチャネルに表示するようにマトリックスを再構成します。