代码之家  ›  专栏  ›  技术社区  ›  diralik

通过2d张量中的值索引pytorch 4d张量

  •  2
  • diralik  · 技术社区  · 6 年前

    • X (A, B, C, D)
    • I 有型 (A, B)

    价值观 [0, C) .


    Y 有型 (A, B, D) ,以便:

    Y[i][j][k] = X[i][j][ I[i][j] ][k]
    
    1 回复  |  直到 6 年前
        1
  •  6
  •   Shai    6 年前

    您可能想使用 torch.gather expand 调适 I 至所需尺寸:

    eI = I[..., None, None].expand(-1, -1, 1, X.size(3))  # make eI the same for the last dimension
    Y = torch.gather(X, dim=2, index=eI).squeeze()
    

    测试代码:

    A = 3 
    B = 4 
    C = 5 
    D = 7
    
    X = torch.rand(A, B, C, D)
    I = torch.randint(0, C, (A, B), dtype=torch.long)
    
    eI = I[..., None, None].expand(-1, -1, 1, X.size(3))
    Y = torch.gather(X, dim=2, index=eI).squeeze()
    
    # manually gather
    refY = torch.empty(A, B, D)
    for i in range(A):
        for j in range(B):
            refY[i, j, :] = X[i, j, I[i,j], :]
    
    (refY == Y).all()
    # Out[]: tensor(1, dtype=torch.uint8)