CNN Predictions

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