PyTorch 中使用 torch.max() 获取最大值索引

这段代码使用 PyTorch 中的 torch.max() 函数获取每个输入的最大值和对应的索引。output 是一个张量，代表模型的输出。torch.max(output, dim=1) 在第一个维度 (dim=1) 上寻找每行的最大值。它返回一个元组，包含每行的最大值和对应的索引。

torch.max(output, dim=1) 返回一个元组 (max_values, indices)，其中 max_values 是一个张量，包含了每行的最大值，indices 是一个张量，包含了每行最大值的索引。[1] 用于取出这个元组中的第二个元素，即最大值的索引。

以下是一个使用 torch.max() 函数的示例：

import torch

# 假设 output 是一个大小为 (batch_size, num_classes) 的张量
output = torch.tensor([[0.1, 0.5, 0.2], [0.3, 0.2, 0.6]])

# 获取每行最大值的索引
pred_class = torch.max(output, dim=1)[1]

print(pred_class)

在这个示例中，我们首先定义了一个大小为 (2, 3) 的张量 output，它代表模型的输出。然后，我们使用 torch.max(output, dim=1)[1] 来获取每行最大值的索引，并将结果存储在 pred_class 变量中。

输出结果如下：

tensor([1, 2])

这表示第一行的最大值索引为 1，第二行的最大值索引为 2。

希望这个示例对你有帮助！如果还有其他问题，请随时提问。