【深度学习】torch.argmax()函数讲解 | pytorch

---

前言

这篇博客也是属于看了好久一直没写，终于写了。

一、两个维度的张量使用torch.argmax()函数

我们直接先举一个例子吧，我们随机生成一个2X3的张量：

import torch
x = torch.rand(2, 3)
print(x)
y0 = torch.argmax(x, dim=0)
print(y0)
y1 = torch.argmax(x, dim=1)
print(y1)
1
2
3
4
5
6
7

结果：

1）dim=0时，返回每一列最大值的索引
2）dim=1时，返回每一行最大值的索引

---

也许看了之后你会说，那这个函数有啥用？我们来看一个实际的例子：
在深度学习的基础内容中，有一篇内容时使用类似Lenet-5这个网络来识别Fashion-Mnist数据集，里面有讲到这个函数：

output = model(b_x)  
output = output.cuda()
pre_lab = torch.argmax(output,1)
1
2
3

之前我们这里设置的batch_size是64，所以将64张输入到网络中，会得到一个64X10的张量，其中64是指有64张图片，10是指10个衣服类别的概率（Fashion—Mnist数据集一共有10个种类的衣服）

在使用torch.argmax(output,1)后，会返回一个长度为64的一维向量，其中64个元素分别代表这64张照片被识别出的某个物品：

二、三个维度的张量使用torch.argmax()函数

我们依然还是举一个简单明了的例子：

import torch

x = torch.rand(3,2,3)
print(x)
1
2
3
4

然后在每一个维度使用这个argmax函数

y0 = torch.argmax(x, dim=0)
print(y0)
y1 = torch.argmax(x, dim=1)
print(y1)
y2 = torch.argmax(x, dim=2)
print(y2)
1
2
3
4
5
6

1）dim=0时，返回每个元素在各个通道最大值的索引号
2）dim=1时，返回n排每一列最大值的索引（n是通道数）
3）dim=2时，返回n排每一排最大值的索引（n是通道数）