PyTorch max()函数取最大值

max() 函数

本人在复现 LeNet5 网络进行 MNIST 数据集的训练，在模型预测阶段，需要用到 max() 函数进行结果的处理。将模型输出的 fc7 层的 10 个结果，变为最后的一个最大概率，也就是具体哪个(0~9)数字的结果。于是，特将 max() 函数的作用记录于此。

导入示例矩阵

import torch
a = torch.Tensor([1,2], [3,5], [6,4])

$$a=\left[\begin{array}{cc}1& 2\\ 3& 5\\ 6& 4\end{array}\right]$$

1. 取全部数据的最大值

>>> pred, idx = a.max()

>>> pred
tensor(6.)
>>> pred.shape
torch.Size([])

$$a=\left[\begin{array}{cc}1& 2\\ 3& 5\\ 6& 4\end{array}\right]$$

2. 按行取最大值，返回最大值和最大值所在索引号

>>> pred, idx = a.max(1)

>>> pred
tensor([2., 5., 6.])
>>> idx
tensor([1, 1, 0])
>>> pred.shape
torch.Size([3])
>>> idx.shape
torch.Size([3])

$$a=\left[\begin{array}{cc}1& 2\\ 3& 5\\ 6& 4\end{array}\right]=\left[\begin{array}{c}2\\ 5\\ 6\end{array}\right]$$

注意，尺寸进行了降维（从二维变为了一维）。

3. 按列取最大值，返回最大值和最大值所在索引号

>>> pred, idx = a.max(0)

>>> pred
tensor([6., 5.])
>>> idx
tensor([2., 1.])
>>> pred.shape
torch.Size([2])
>>> idx.shape
torch.Size([2])

$$a=\left[\begin{array}{cc}1& 2\\ 3& 5\\ 6& 4\end{array}\right]=\left[\begin{array}{cc}6& 5\end{array}\right]$$

注意，尺寸进行了降维（从二维变为了一维）。