问题描述与目标
在 pytorch 张量操作中，我们经常会遇到这样的需求：给定一个主张量 a，以及一个或多个参考张量（例如 b 和 c），需要生成一个与 a 形状相同的布尔掩码。这个掩码的每个位置为 true 当且仅当 a 中对应位置的元素存在于任何一个参考张量中。例如，如果 a = [1, 234, 54, 6543, 55, 776]，b = [234, 54]，c = [55, 776]，我们期望得到的掩码是 [false, true, true, false, true, true]。
低效的循环实现方法
一种直观但效率较低的实现方式是遍历参考张量中的每个元素，然后使用相等性比较和求和操作来构建掩码。这种方法涉及隐式或显式的循环，对于大型张量而言，其计算成本会迅速增加。
以下是这种方法的示例代码：
import torch
# 定义主张量和参考张量
a = torch.tensor([1, 234, 54, 6543, 55, 776])
b = torch.tensor([234, 54])
c = torch.tensor([55, 776])
# 使用循环和求和构建掩码
# 对于b中的每个元素i，检查a中哪些元素等于i，得到一个布尔张量
# 然后将这些布尔张量求和，再转换为布尔类型
a_masked_b = sum(a == i for i in b).bool()
a_masked_c = sum(a == i for i in c).bool()
# 将来自b和c的掩码进行逻辑或操作
a_masked = a_masked_b + a_masked_c # 或者使用 a_masked_b | a_masked_c
print(f"主张量 a: {a}")
print(f"参考张量 b: {b}")
print(f"参考张量 c: {c}")
print(f"通过循环方法生成的掩码: {a_masked}")
# 预期输出: tensor([False, True, True, False, True, True])注意事项： 这种方法虽然能够得到正确结果，但其性能瓶颈在于内部的循环操作。对于包含大量元素的张量或多个参考张量的情况，这种方法会非常慢，不推荐在生产环境中使用。
高效的 PyTorch 解决方案：torch.isin
PyTorch 提供了 torch.isin 函数，专门用于检查一个张量中的元素是否包含在另一个张量中。这个函数在底层进行了高度优化，通常比手动循环快数倍甚至数十倍。
torch.isin(elements, test_elements) 函数接受两个参数：
- elements: 需要被检查的张量，即我们的主张量 a。
- test_elements: 包含用于测试的元素集合的张量，即所有参考张量合并后的集合。
为了将多个参考张量（如 b 和 c）合并成一个用于测试的集合，我们可以使用 torch.cat() 函数将它们拼接起来。
以下是使用 torch.isin 的示例代码：
import torch
# 定义主张量和参考张量
a = torch.tensor([1, 234, 54, 6543, 55, 776])
b = torch.tensor([234, 54])
c = torch.tensor([55, 776])
# 将所有参考张量合并成一个测试集合
all_test_elements = torch.cat([b, c])
# 使用 torch.isin 生成掩码
a_masked_isin = torch.isin(a, all_test_elements)
print(f"主张量 a: {a}")
print(f"合并后的测试元素集合: {all_test_elements}")
print(f"通过 torch.isin 生成的掩码: {a_masked_isin}")
# 预期输出: tensor([False, True, True, False, True, True])优势： torch.isin 函数的底层实现通常利用了哈希表或排序等高效算法，能够以远超显式循环的速度完成元素包含性检查。这是处理大规模张量时推荐的方法。
总结
在 PyTorch 中检查一个张量中的元素是否包含在其他张量中，并生成相应的布尔掩码，最推荐且高效的方法是使用 torch.isin 函数。通过将所有参考张量合并成一个单一的测试集合，torch.isin 能够以优化的方式完成此任务，避免了低效的 Python 循环，从而显著提升代码性能和可读性。在实际应用中，尤其是在处理大型数据集时，始终优先考虑使用 PyTorch 提供的向量化操作和优化函数，如 torch.isin。
