Understanding training

[ptrblck]

@dmandair

One follow up to make sure I’ve got this - is the use of argmax here distinct from the use of torch.max in the Pytorch tutorial (ie in the line: _, preds = torch.max(outputs, 1), see the entire training epoch below).
[…]
the Pytorch tutorial’s use of torch.max should extend to the multi-class classification case as well, correct? Since torch.max returns a tuple of the max and indices for those values, and I believe that function only uses the indices themselves, I think it should work with multi-class labels but definitely let me know if I’m mistaken on that.

Yes, you are exactly right. torch.max would return the values with their corresponding indices, while torch.argmax only returns the indices which correspond to the max. values.
As you can see in the tutorial, the values are not used _, preds = torch.max(outputs, 1) to it would be equal to preds = torch.argmax(outputs, 1).

@dmandair

but never detaches the output in torch.max. Wouldn’t that similarly store the computation graph or does torch.max behave differently?

That’s also a great question! You are generally right and should be careful about detaching the tensors. However, since the returned indices from torch.argmax are not differentiable, the computation graph won’t be stored. If you are using torch.max, note that the values will be attached to the computation graph while the indices also won’t be attached to it. You can check it via:

output = torch.randn(10, 10, requires_grad=True)

val, idx = torch.max(output, dim=1)

# a valid grad_fn shows that this tensor is attached to the computation graph
print(val.grad_fn)
> <MaxBackward0 object at 0x7fe084afbbb0> 

print(idx.grad_fn)
> None

idx = torch.argmax(output, dim=1)
print(idx.grad_fn)
> None