关于torch.argmax()的在深度学习中应用的理解

今天在debug过程中看到代码中使用了torch.argmax()这个函数。之前在用的时候只是知道返回最大值对应的位置索引，但没有深入思考其含义。

在查了大量资料之后，发现很多博客仅仅是照搬了官方文档中的释义，并没有对该函数的使用场景及含义做过多的解释，导致很多初学者在查阅这一函数的用法时一头雾水。

下面就结合自己的一些理解对这个函数进行一个直观的解释。

---

在深度学习中，torch.argmax()这个函数最常见的用处就是在图像分类任务下，筛选出网络预测值Prediction中最大概率值对应的class ID。

在Pytorch框架中，一个Batch的张量是以 [N, C, H, W] 格式进行存储的，其中N为Batch size，C为通道数Channel(Number of Class)，H和W为图像的高度和宽度。

以最简单的二分类为例，网络的输出是一个 [N, 2, H, W] 的张量。

此时调用torch.argmax(prediction, dim=1)，其具体含义为：对于N组预测结果图，遍历每一组中的每一个像素，找到两张预测图像素值的最大值对应的channel。具体过程如下图所示：

对于一个[N, 2, H, W]的input，最终torch.argmax(predition, dim=1)的输出是一个[N, H, W]的张量，也就是一组包含有batch size个预测结果图的y'。

此时再将y'和标签label输入如交叉熵损失函是一类的loss function就可以进行loss值的计算了。