在使用 pytorch 进行训练时,会使用使用到改行代码:
predict = torch.max(outputs.data, 1)[1]
其中 output 为模型的输出,该函数主要用来求 tensor 的最大值。

每次看到都不太理解 torch.max() 的使用,为了下次看到或者写道时不会忘记,特意详细了解其用法。
torch.max(input:tensor, dim:index)
该函数有两个输入:
inputs: tensor,第一个参数为一个张量
dim: index,第二个参数为一个整数[-2-1],dim=0表示计算每列的最大值,dim=1表示每行的最大值
如:

import torch

output = torch.tensor([[1, 23], [3, 45]])
predict = torch.max(output, dim=0)
print(predict)
predict = torch.max(output, dim=1)
print(predict)

输出为:
torch.return_types.max(
values=tensor([3, 4, 5]),
indices=tensor([1, 1, 1]))

torch.return_types.max(
values=tensor([3, 5]),
indices=tensor([2, 2]))

输出返回最大值以及对应的索引值。因此,predict = torch.max(outputs.data, 1)[1] 是计算模型中每个类别的最大值并返回其索引值,即该类别的标签值。

Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐