Reshaping Pytorch tensor

1
votes

I have a tensor of size (24, 2, 224, 224) in Pytorch.

  • 24 = batch size
  • 2 = matrixes representing foreground and background
  • 224 = image height dimension
  • 224 = image width dimension

This is the output of a CNN that performs binary segmentation. In each cell of the 2 matrixes is stored the probability for that pixel to be foreground or background: [n][0][h][w] + [n][1][h][w] = 1 for every coordinate

I want to reshape it into a tensor of size (24, 1, 224, 224). The values in the new layer should be 0 or 1 according to the matrix in which the probability was higher.

How can I do that? Which function should I use?

1
pythonarraysnumpymatrixpytorch

1 Answers

2
votes

Using torch.argmax() (for PyTorch +0.4):

prediction = torch.argmax(tensor, dim=1) # with 'dim' the considered dimension 
prediction = prediction.unsqueeze(1) # to reshape from (24, 224, 224) to (24, 1, 224, 224)

If the PyTorch version is below 0.4.0, one can use tensor.max() which returns both the max values and their indices (but which isn't differentiable over the index values):

_, prediction = tensor.max(dim=1)
prediction = prediction.unsqueeze(1) # to reshape from (24, 224, 224) to (24, 1, 224, 224)