Hi everyone,
In this example:
for images, labels in test_loader:
images = Variable(images).cuda()
outputs = cnn(images)
_, predicted = torch.max(outputs.data, 1)
I would like to know why we define the predictions by finding the maximum between outputs.data
and 1
.
I did a Google search and it seems like you got that code from this example: https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/04%20-%20Convolutional%20Neural%20Network/main-gpu.py
The CNN network is trained with CrossEntropyLoss
which will mean the maximum value given by the final layer will be the predicted result. The 1
given as a parameter to torch.max
is simply the dimension in outputs.data
we want to check for the maximum:
torch.max(input, dim, max=None, max_indices=None) -> (Tensor, LongTensor)
Source: http://pytorch.org/docs/torch.html#torch.max
1 Like