CODEWITHSUNDEEP
pythonarraysnumpymatrixpytorch
Jimbo
Jimbo

Reputation: 494

Reshaping Pytorch tensor

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

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?

Upvotes: 1

Views: 1906

Answers (1)

benjaminplanche
benjaminplanche

Reputation: 15139

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)

Upvotes: 2

Related Questions