PyTorch向量交集快速检测方法探究
2026-6-19
快速判断两个PyTorch向量是否存在共同元素的最优方案
问题
给定两个PyTorch向量v1、v2，如何以最快速度判断二者是否存在共同元素？
环境设置
- 数据类型：
torch.int64 - v1：长度106-108，预计算的唯一有序元素向量
- v2：长度105-107，每次都会变化
- 典型场景：二者无共同元素
- 优先支持CPU和GPU双平台
已验证的方案
torch.isin(v1, v2, assume_unique=True)：设置assume_unique=True可使速度提升高达10倍，必须利用这个参数优化torch.searchsorted匹配验证：先通过searchsorted在有序v1中定位v2元素的位置，再验证位置对应元素是否与v2相等，多数场景下速度最快；但后续的求和/存在性判断步骤耗时占比高，仍有优化空间- 其他方案：性能更差，不推荐
CPU性能测试结果
| 方法 | 耗时 | 设备 | L1长度 | L2长度 | 交集大小 | L1唯一元素数 | L2唯一元素数 |
|---|---|---|---|---|---|---|---|
| isin direct | 1.365 | cpu | 1,000,000 | 1,000,000 | 0 | 1,000,000 | 1,000,000 |
| isin assume_unique | 0.778 | cpu | 1,000,000 | 1,000,000 | 0 | 1,000,000 | 1,000,000 |
| search sorted | 0.136 | cpu | 1,000,000 | 1,000,000 | 0 | 1,000,000 | 1,000,000 |
| isin direct | 13.286 | cpu | 10,000,000 | 1,000,000 | 0 | 10,000,000 | 1,000,000 |
| isin assume_unique | 1.823 | cpu | 10,000,000 | 1,000,000 | 0 | 10,000,000 | 1,000,000 |
| search sorted | 0.321 | cpu | 10,000,000 | 1,000,000 | 0 | 10,000,000 | 1,000,000 |
GPU性能测试结果
注：测试中torch.randint存在bug导致交集大小数据不准确，但耗时数据有效；该bug已在后续版本修复。
| 方法 | 耗时 | 设备 | L1长度 | L2长度 | 交集大小 | L1唯一元素数 | L2唯一元素数 |
|---|---|---|---|---|---|---|---|
| isin direct | 0.005 | cuda | 1,000,000 | 1,000,000 | 311,872 | 1,000,000 | 1,000,000 |
| isin assume_unique | 0.003 | cuda | 1,000,000 | 1,000,000 | 311,872 | 1,000,000 | 1,000,000 |
| search sorted | 0.001 | cuda | 1,000,000 | 1,000,000 | 311,872 | 1,000,000 | 1,000,000 |
| isin direct | 0.034 | cuda | 10,000,000 | 1,000,000 | 1,000,000 | 10,000,000 | 1,000,000 |
| isin assume_unique | 0.011 | cuda | 10,000,000 | 1,000,000 | 1,000,000 | 10,000,000 | 1,000,000 |
| search sorted | 0.003 | cuda | 10,000,000 | 1,000,000 | 99,480 | 10,000,000 | 1,000,000 |
| isin direct | 0.296 | cuda | 100000000 | 1000000 | 1023183.0 | 98845253.0 | 999877.0 |
| isin assume_unique | 0.081 | cuda | 100000000 | 1000000 | 999877.0 | 98845253.0 | 999877.0 |
| search sorted | 0.003 | cuda | 100000000 | 1000000 | 10019.0 | 98845253.0 | 999877.0 |
特殊场景测试（向量规模相近）
当两个向量规模相近（如均为10^7）时，各方法性能差异极小：
场景1：双10^7规模向量
| 方法 | 耗时 | 设备 | L1长度 | L2长度 | 交集大小 | L1唯一元素数 | L2唯一元素数 |
|---|---|---|---|---|---|---|---|
| isin direct | 0.041 | cuda | 10000000 | 10000000 | 4965335 | 9988353 | 9988316 |
| isin direct v2→v1 | 0.034 | cuda | 10000000 | 10000000 | 4965266 | 9988353 | 9988316 |
| isin assume_unique | 0.020 | cuda | 10000000 | 10000000 | 4956632 | 9988353 | 9988316 |
| isin assume_unique v2→v1 | 0.020 | cuda | 10000000 | 10000000 | 4956632 | 9988353 | 9988316 |
| search sorted | 0.026 | cuda | 10000000 | 10000000 | 4965266 | 9988353 | 9988316 |
场景2：10^8 vs 5×10^7规模向量
| 方法 | 耗时 | 设备 | L1长度 | L2长度 | 交集大小 | L1唯一元素数 | L2唯一元素数 |
|---|---|---|---|---|---|---|---|
| isin direct | 0.354 | cuda | 100000000 | 50000000 | 50574973 | 98845711 | 49711185 |
| isin direct v2→v1 | 0.229 | cuda | 100000000 | 50000000 | 49996126 | 98845711 | 49711185 |
| isin assume_unique | 0.138 | cuda | 100000000 | 50000000 | 49707311 | 98845711 | 49711185 |
| isin assume_unique v2→v1 | 0.137 | cuda | 100000000 | 50000000 | 49707311 | 98845711 | 49711185 |
| search sorted | 0.184 | cuda | 100000000 | 50000000 | 49996126 | 98845711 | 49711185 |
实验代码示例
方案1（基础isin）
m = torch.isin(v1, v2) has_common = m.sum().item() > 0
方案2（带unique优化的isin）
v2unique = torch.unique(v2) m = torch.isin(v1, v2unique, assume_unique=True) has_common = m.sum().item() > 0
方案3（searchsorted验证）
ts = torch.searchsorted(v1, v2) # 处理超出v1范围的索引 ts[ts >= len(v1)] = len(v1) - 1 # 直接判断是否存在匹配，无需求和 has_common = (v1[ts] == v2).any().item()
内容的提问来源于stack exchange，提问作者Alexander Chervov
相关产品推荐
相关产品推荐
