pytorch: torch.max() 使用与理解
·
在使用 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, 2, 3], [3, 4, 5]])
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] 是计算模型中每个类别的最大值并返回其索引值,即该类别的标签值。
更多推荐
已为社区贡献11条内容
所有评论(0)