如何在PyTorch张量的指定索引子集上求解argmax/argmin

解决方案 #

完全可以通过向量化操作实现，无需循环，所有操作会保留原有设备属性，不需要额外数据迁移。

实现逻辑 #

1. 用指定的索引子集从原distance tensor中提取对应位置的数值
2. 对提取出的子集求argmax，得到最大值在子集中的相对位置
3. 用这个相对位置从索引子集里取出原张量对应的原始索引即可

代码示例（PyTorch） #

import torch

# 示例输入
distance = torch.tensor([5, 10, 2, 3, 4], device='cuda:0')
indices = torch.tensor([0, 2, 3], device='cuda:0')

# 核心操作
subset_values = distance[indices]
subset_argmax = subset_values.argmax()
# 加unsqueeze(0)是为了匹配你需要的[tensor([0], device='cuda:0')] 一维张量输出格式
result = indices[subset_argmax].unsqueeze(0)

print(result)
# 输出：tensor([0], device='cuda:0')

补充说明 #

如果不需要保留一维张量格式，直接取indices[subset_argmax]就会得到标量张量tensor(0, device='cuda:0')，可以根据实际需求调整是否添加unsqueeze操作。

内容的提问来源于stack exchange，提问作者Vojtooo