如何对二维Tensor逐元素应用自定义判断函数？

直接使用PyTorch原生的isin方法即可实现高效的批量判断，代码如下：

import torch
t = torch.tensor([[2,3],[4,6]])
overlap = [2, 6]

result = t.isin(overlap)
print(result)

输出结果：

tensor([[ True, False],
        [False,  True]])

方案优势 #

• isin是PyTorch底层优化的矢量化操作，基于C++实现，会对整个Tensor进行批量处理，彻底避免了Python层面逐元素循环的开销。
• 对比你原先用lambda逐元素判断的方式，这个方案在处理大规模Tensor和大列表时性能提升极其显著——后者每个元素都要触发Python解释器调用，数据量越大，效率差距越明显。

额外说明 #

如果overlap本身是PyTorch Tensor，也可以直接传入isin，效果完全一致，无需额外转换。

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