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]$$

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

总结

以上是生活随笔为你收集整理的PyTorch max()函数取最大值的全部内容，希望文章能够帮你解决所遇到的问题。

如果觉得生活随笔网站内容还不错，欢迎将生活随笔推荐给好友。