Suppose I have a tensor like [A,A,A,A...,A].
How can I quickly obtain [[A],[A],[A],...,[A]] as a tensor in torch?
Suppose I have a tensor like [A,A,A,A...,A].
How can I quickly obtain [[A],[A],[A],...,[A]] as a tensor in torch?
You can use torch.chunk as the inverse of cat, but it looks like you want unsqueeze(1):
A = torch.randn(2, 3)
A_rep = (A, A, A, A, A, A, A, A)
catted = torch.cat(A_rep)
#uncatted = torch.chunk(catted, len(A_rep))
catted.unsqueeze(1)
From the doc, use split or reshape.
>>> a = torch.arange(10).reshape(5,2)
>>> a
tensor([[0, 1],
[2, 3],
[4, 5],
[6, 7],
[8, 9]])
>>> torch.split(a, 2)
(tensor([[0, 1],
[2, 3]]),
tensor([[4, 5],
[6, 7]]),
tensor([[8, 9]]))
>>> torch.split(a, [1,4])
(tensor([[0, 1]]),
tensor([[2, 3],
[4, 5],
[6, 7],
[8, 9]]))
You are looking for torch.unbind which does exactly what you want.
import torch
tensor = torch.rand(3, 4, 5) # tensor of shape (3, 4, 5)
l = tensor.unbind(dim=0) # list of 3 tensors of shape (4, 5)
You can unbind along whatever dimension you want, default is dim 0. You can use python destructuring to quickly get the tensors in variables:
a, b, c = tensor.unbind(dim=0) # 3 tensors of shape (4, 5)