Best way to convert a list to a tensor?

9

If someone is looking into the performance aspects of this, I’ve done a small experiment. In my case, I needed to convert a list of scalar tensors into a single tensor.

import torch
torch.__version__ # 1.10.2
x = [torch.randn(1) for _ in range(10000)]
torch.cat(x).shape, torch.stack(x).shape # torch.Size([10000]), torch.Size([10000, 1])

%timeit torch.cat(x) # 1.5 ms ± 476 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit torch.cat(x).reshape(-1,1) # 1.95 ms ± 534 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit torch.stack(x) # 5.36 ms ± 643 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

My conclusion is that even if you want to have the additional dimension of torch.stack, using torch.cat and then reshape is better.

1 Like