인덱싱 한 PyTorch 텐서는

Question

나는 이런 식으로 뭔가를 보이는 3 × 3 PyTorch LongTensor이 index_select 사용

환언

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] 

매트릭스 :

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

내 의도 된 출력은 아래의 2 × 3 × 3이다 FloatTensor 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 x 2 행렬을 색인 할 때)를 반영합니다.

마지막으로 __.permute(2,0,1)은 출력이 각 열 대신에 B의 각 차원을 별도의 채널에서 볼 수 있도록 행렬을 바꿉니다.