Add extra dimension to an axes
I have a batch of segmentation masks of shape [5,1,100,100] (batch_size x dims x ht x wd) which I have to display in tensorboardX with an RGB image batch [5,3,100,100]. I want to add two dummy dimensions to the second axes of the segmentation mask to make it [5,3,100,100] so there will not be any dimension mismatch error when I pass it to torch.utils.make_grid. I have tried unsqueeze, expand and view but I am not able to do it. Any suggestions?

You can use expand, repeat, or repeat_interleave:

import torch

x = torch.randn((5, 1, 100, 100))
x1_3channels = x.expand(-1, 3, -1, -1)
x2_3channels = x.repeat(1, 3, 1, 1)
x3_3channels = x.repeat_interleave(3, dim=1)

print(x1_3channels.shape)  # torch.Size([5, 3, 100, 100])
print(x2_3channels.shape)  # torch.Size([5, 3, 100, 100])
print(x3_3channels.shape)  # torch.Size([5, 3, 100, 100])


Note that, as stated in the docs:


torch.expand():



  Expanding a tensor does not allocate new memory, but only creates a new view on the existing tensor where a dimension of size one is expanded to a larger size by setting the stride to 0. Any dimension of size 1 can be expanded to an arbitrary value without allocating new memory.



torch.repeat():



  Unlike expand(), this function copies the tensor’s data.

