因為argmax函數不能求導,也就不會有梯度
backward的反向傳播的梯度到loss是終止了
其實也很容易看出這個問題,如果最后一步使用了argmax,loss的屬性 requires_grad=False,直接調用backward會提示你需要給loss增加梯度屬性,但這“治表不治本 ”
網上有兩個ref參考
https://discuss.pytorch.org/t/torch-argmax-cause-loss-backward-dont-work/64782/2
https://discuss.pytorch.org/t/backpropagation-issue-when-using-argmax/100335/6
以及 TORCH.TENSOR.GRAD 的文檔
This attribute is None by default and becomes a Tensor the first time a call to backward() computes gradients for self. The attribute will then contain the gradients computed and future calls to backward() will accumulate (add) gradients into it.
這段話說明了梯度的兩個特點:一是第一次調用backward才會計算梯度,之前是None,二是梯度是可以累加的,如果沒有手動清掉的話
補充:
知乎直接給我推了這個問題...: 怎么克服神經網絡訓練中argmax的不可導性
不過通常來說還是loss設計有問題,比如我本該用回歸loss用成了分類loss,才被迫使用argmax