Reshaping tensor

1

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 sholud be 0 or 1 according to the matrix in which the probability was higher.

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

2

You could just use torch.argmax():

a = torch.randn(1, 2, 4, 4)
print(a)
b = torch.argmax(a, 1)
b.unsqueeze_(1)
print(b)
3

Running your code I get the following error:

AttributeError: module ‘torch’ has no attribute ‘argmax’

That’s very strange, I’ve used torch.argmax before. Something obvious that I’m missing?

4

Which PyTorch version are you using?
As far as I know, torch.argmax was introduced in 0.4.0.
If you have an older version, you could update to the latest stable release following these instructions or use:

_, idx = torch.max(a, 1)
5

That’s it!
I was using torch 0.4.0 some time ago, then I rolled back to 0.3.1 to conform with my work group.

Thank you very much! You’re as helpful as usual.

1 Like