close

How to use torch.stack function

Hello Guys, How are you all? Hope You all Are Fine. Today We Are Going To learn about How to use torch.stack function in Python. So Here I am Explain to you all the possible Methods here.

Without wasting your time, Let’s start This Article.

Table of Contents

How to use torch.stack function?

  1. How to use torch.stack function?

    Stacking requires same number of dimensions. One way would be to unsqueeze and stack. For example:

  2. use torch.stack function

    Stacking requires same number of dimensions. One way would be to unsqueeze and stack. For example:

Method 1

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

Method 2

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.

Summery

It’s all About this issue. Hope all Methods helped you a lot. Comment below Your thoughts and your queries. Also, Comment below which Method worked for you? Thank You.

Also, Read