PyTorch中依据张量b元素是否在c中更新张量a的最优方法问询

你的实现思路可行，但存在关键性能短板：将张量转换为Python集合、列表的操作会脱离PyTorch的张量计算体系，在数据量较大或使用GPU加速的场景下，会产生额外的转换开销，拖慢整体运行速度。

更优的实现方式是直接使用PyTorch原生的torch.isin函数，完全基于张量操作完成需求，代码更简洁且性能更高：

import torch

a = torch.tensor([1.0, 2.0, 3.0, 4.0])
b = torch.tensor([0, 2, 1, 3])
c = torch.tensor([1, 2])

# 直接生成布尔掩码，无需转换为Python对象
update_mask = torch.isin(b, c)
a[update_mask] *= 2

为什么这个方法更优？ #

• 无额外转换开销：torch.isin直接在张量层面完成成员检测，不需要把c转成Python集合、把b转成列表，避免了张量与Python对象之间的数据拷贝和类型转换。
• 支持GPU加速：如果你的张量在GPU上，torch.isin会自动利用CUDA并行计算，而原方法必须把数据拉回CPU处理，性能差距会非常大。
• 代码更简洁：省去了手动生成掩码的列表推导式，逻辑更直观。

即使c中的元素存在重复，torch.isin也能正确处理（只要元素存在就标记为True），完全覆盖你的需求场景。

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