PyTorch中匹配多个布尔条件更新张量值的最优方法

PyTorch 原生的==比较运算符不支持直接传入目标值列表实现多值或逻辑匹配。
针对需要匹配大量目标值的场景，最佳实现方案是使用PyTorch内置的torch.isin()函数，该函数专门用于判断张量中的元素是否属于给定的目标值集合，底层自动等效实现所有目标值的OR匹配逻辑，无需手动拼接大量|连接的判断条件。

对应你的示例，改写后的代码如下：

import torch

my_tensor = torch.tensor([0, 1, 2, 3, 4, 5])
# 所有需要匹配的目标值统一放在列表中维护即可
match_targets = [1, 4, 5]
condition = torch.isin(my_tensor, torch.tensor(match_targets))
my_tensor[condition] = 0

print(my_tensor)
# 输出: tensor([0, 0, 2, 3, 0, 0])

这个方案的优势：

• 代码可维护性强，无论需要匹配多少个目标值，只需要修改match_targets列表即可，不需要重复编写my_tensor==x的判断片段
• 执行效率高，torch.isin底层做了专门的性能优化，比手动拼接大量OR条件的运行速度更快，在大张量、多匹配值的场景下性能优势尤其明显
• 适配任意维度的张量输入，不需要针对张量形状额外调整逻辑

如果你使用的是1.10之前的老旧PyTorch版本（该版本才正式加入torch.isin），可以用广播机制实现等效逻辑，注意该方案在张量和目标值列表规模都很大时内存占用会高于官方内置函数：

match_targets = torch.tensor([1, 4, 5])
# 利用广播逐元素比较后，沿目标值维度做any判断，等效OR逻辑
condition = (my_tensor.unsqueeze(-1) == match_targets).any(dim=-1)

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