If I understand correctly what you're trying to do is stack the outputted mini-batches together into a single batch. My bet is that your last batch is partially filled (only has 16 elements instead of 32).
Instead of using torch.stack (creating a new axis), I would simply concatenate with torch.cat on the batch axis (axis=0). Assuming matrices is a list of torch.Tensors.
torch.cat(matrices).cpu().detach().numpy()
As torch.cat concatenates on axis=0 by default.
When we have tensors that differ in size only on the first dimension, as of PyTorch v1.7.0, we can use torch.vstack() to stack it along axis 0. Using torch.stack() fails here because it expects all the tensors to be of same shape.
Here is a reproducible illustration matching your problem description:
# sample tensors (as per your size)
In [65]: t1 = torch.randn([32, 1, 600, 600])
In [66]: t2 = torch.randn([16, 1, 600, 600])
# vertical stacking (i.e., stacking along axis 0)
In [67]: stacked = torch.vstack([t1, t2])
# check shape of output
In [68]: stacked.shape
Out[68]: torch.Size([48, 1, 600, 600])
we get 48 (32 + 16) as the size of first dimension in the result because we're stacking tensors along that dimension.
Note:
You can also initialize the result tensor, say stacked, by explicitly calculating the shape and pass this tensor as a parameter to out= kwarg of torch.vstack() if you want to write the result to a specific tensor, for instance updating the values of existing tensor (of same shape). However, this is optional.
# calculate new shape of stacking
In [80]: newshape = (t1.shape[0] + t2.shape[0], *t1.shape[1:])
# allocate an empty tensor, filled with garbage values
In [81]: stacked = torch.empty(newshape)
# stack it along axis 0 and write the result to `stacked`
In [83]: torch.vstack([t1, t2], out=stacked)
# check shape/size
In [84]: stacked.shape
Out[84]: torch.Size([48, 1, 600, 600])
Imagine have n tensors. If we stay in 3D, those correspond to volumes, namely rectangular cuboids. Stacking corresponds to combining those n volumes on an additional dimension: here a 4th dimension is added to host the n 3D volumes. This operation is in clear contrast with concatenation, where the volumes would be combined on one of the existing dimensions. So concatenation of three-dimensional tensors would result in a 3D tensor.
Here is a possible representation of the stacking operations for limited dimensions sizes (up to three-dimensional inputs):
Where you chose to perform the stacking defines along which new dimension the stack will take place. In the above examples, the newly created dimension is last, hence the idea of "added dimension" makes more sense.
In the following visualization, we observe how tensors can be stacked on different axes. This in turn affects the resulting tensor shape
For the 1D case, for instance, it can also happen on the first axis, see below:
With code:
>>> x_1d = list(torch.empty(3, 2)) # 3 lines >>> torch.stack(x_1d, 0).shape # axis=0 stacking torch.Size([3, 2]) >>> torch.stack(x_1d, 1).shape # axis=1 stacking torch.Size([2, 3])Similarly for two-dimensional inputs:
With code:
>>> x_2d = list(torch.empty(3, 2, 2)) # 3 2x2-squares >>> torch.stack(x_2d, 0).shape # axis=0 stacking torch.Size([3, 2, 2]) >>> torch.stack(x_2d, 1).shape # axis=1 stacking torch.Size([2, 3, 2]) >>> torch.stack(x_2d, 2).shape # axis=2 stacking torch.Size([2, 2, 3])
With this state of mind, you can intuitively extend the operation to n-dimensional tensors.
Very simple! I will use 4 varibles for this example. The function torch.stack is very similar to numpy (vstack and hstack). Example :
t1 = torch.tensor([1,2,3])
t2 = torch.tensor([4,5,6])
t3 = torch.tensor([7,8,9])
t4 = torch.tensor([10,11,12])
if you try this cmd
>> torch.stack((t1,t2,t3,t4),dim=1).size()
>> torch.Size([3, 4])
if you change dim=1 with dim=0
>> torch.stack((t1,t2,t3,t4),dim=0).size()
>> torch.Size([4, 3])
In first case you have a tensor with 3x4 dimensional but in last case you have a 4x3 tensor. Test this code without using .size() !
>> torch.stack((t1,t2,t3,t4),dim=1)
>> tensor([[ 1, 4, 7, 10],
[ 2, 5, 8, 11],
[ 3, 6, 9, 12]])
>> torch.stack((t1,t2,t3,t4),dim=0)
>> tensor([[ 1, 2, 3],
[ 4, 5, 6],
[ 7, 8, 9],
[10, 11, 12]])
Happy Coding!
Stacking requires same number of dimensions. One way would be to unsqueeze and stack. For example:
a.size() # 2, 3, 4
b.size() # 2, 3
b = torch.unsqueeze(b, dim=2) # 2, 3, 1
# torch.unsqueeze(b, dim=-1) does the same thing
torch.stack([a, b], dim=2) # 2, 3, 5
Using pytorch 1.2 or 1.4 arjoonn's answer did not work for me.
Instead of torch.stack I have used torch.cat with pytorch 1.2 and 1.4:
>>> import torch
>>> a = torch.randn([2, 3, 4])
>>> b = torch.randn([2, 3])
>>> b = b.unsqueeze(dim=2)
>>> b.shape
torch.Size([2, 3, 1])
>>> torch.cat([a, b], dim=2).shape
torch.Size([2, 3, 5])
If you want to use torch.stack the dimensions of the tensors have to be the same:
>>> a = torch.randn([2, 3, 4])
>>> b = torch.randn([2, 3, 4])
>>> torch.stack([a, b]).shape
torch.Size([2, 2, 3, 4])
Here is another example:
>>> t = torch.tensor([1, 1, 2])
>>> stacked = torch.stack([t, t, t], dim=0)
>>> t.shape, stacked.shape, stacked
(torch.Size([3]),
torch.Size([3, 3]),
tensor([[1, 1, 2],
[1, 1, 2],
[1, 1, 2]]))
With stack you have the dim parameter which lets you specify on which dimension you stack the tensors with equal dimensions.