[Python] How to map values from a 3D tensor to a 1D tensor in PyTorch?

Discussão em 'Python' iniciado por Stack, Janeiro 26, 2025.

  1. Stack

    Stack Membro Participativo

    I'm stuck with a Pytorch problem and could use some help:

    I've got two tensors:

    • A 3D tensor (shape: i, j, j) with integer values from 0 to n
    • A 1D tensor (shape: n)

    I need to create a new tensor that's the same shape as the first one (i, j, j), but where each value is replaced by the corresponding value from the second tensor. Like using the values in the first tensor as indices for the second one.

    Here's a quick example:

    import torch

    # My tensors
    big_tensor = torch.randint(0, 256, (10, 25, 25))
    small_tensor = torch.rand(256)

    # What I'm trying to do
    result = magic_function(big_tensor, small_tensor)

    # How it should work
    print(big_tensor[0, 0, 0]) # Let's say this outputs 42
    print(small_tensor[42]) # This might output 0.7853
    print(result[0, 0, 0]) # This should also be 0.7853


    I'm looking for a performant way to do this, preferably not using loops, as both tensors can be quite big. Is there an efficient Pytorch operation or method I can use for this?

    Any help would be awesome - thanks in advance!

    Continue reading...

Compartilhe esta Página