本文介绍如何将逐列遍历的pytorch模式筛选函数完全向量化，通过空间维度展平与布尔掩码逻辑组合，在毫秒级完成对51×265×23×23张量的“存活模式”（含-1/0/1三态）快速过滤，彻底消除python for循环。
本文介绍如何将逐列遍历的pytorch模式筛选函数完全向量化，通过空间维度展平与布尔掩码逻辑组合，在毫秒级完成对51×265×23×23张量的“存活模式”（含-1/0/1三态）快速过滤，彻底消除python for循环。
在处理高维时序张量（如 torch.Size([51, 265, 23, 23])）时，常需按时间步（首维）和模式通道（第二维）进行条件筛选。本例中，“存活模式”定义为：在最后时间步（即 tensor_sims[-1]）中，该模式的23×23空间区域同时包含 -1、0、1 三种状态值；其余情况视为“死亡”，需被剔除。
原始实现使用列表推导式配合 torch.unique().numel() == 3 判断，虽逻辑清晰，但因显式 Python 循环 + 每次调用 unique()（内部含排序与去重），导致严重性能瓶颈——尤其当 tensor_sims.shape[1] = 265 时，开销呈线性增长。
✅ 优化核心思路：
避免逐列迭代与重复 unique 计算，转而利用 PyTorch 张量广播与布尔逻辑的并行能力：
- 聚焦关键切片：仅提取最后一帧 tensor_sims[-1]（形状 [265, 23, 23]）；
- 空间展平：flatten(1) 将后两维压成一维，得到 [265, 529] 张量，每行对应一个模式的所有像素；
- 三态存在性检测：分别生成布尔矩阵 x_ == -1、x_ == 0、x_ == 1（各为 [265, 529]），再沿列方向（dim=1）取 .any()，得到三个长度为 265 的布尔向量，表示每个模式是否至少含对应值；
- 逻辑组合：用 .logical_and() 连接三者，仅当三态全部存在时结果为 True；
- 掩码索引：直接用于高级索引 tensor_sims[:, mask]，一步完成通道裁剪。
以下是完整优化代码：
def filter_patterns(tensor_sims):
# 提取最后时间步，并展平空间维度 (265, 23, 23) → (265, 529)
x_last = tensor_sims[-1].flatten(1)
# 分别检测每种状态是否存在（沿空间维度 any）
has_neg1 = (x_last == -1).any(dim=1) # shape: [265]
has_zero = (x_last == 0).any(dim=1) # shape: [265]
has_pos1 = (x_last == 1).any(dim=1) # shape: [265]
# 三态必须同时存在 → 逻辑与
mask = has_neg1 & has_zero & has_pos1 # shape: [265], dtype: bool
# 向量化索引：保留满足条件的模式通道
tensor_sims = tensor_sims[:, mask]
print(f'Number of patterns: {tensor_sims.shape[1]}')
return tensor_sims
⚠️ 关键注意事项：
- flatten(1) 等价于 reshape(x_last.size(0), -1)，确保空间信息不丢失且维度对齐；
- 必须使用 .any(dim=1) 而非 .all() 或 unique——前者复杂度 O(N)，后者为 O(N log N)，且支持 GPU 加速；
- 掩码 mask 为布尔张量，PyTorch 高级索引自动广播，无需 torch.nonzero() 转换索引；
- 若输入张量在 CUDA 上，该函数全程保持 GPU 原生运算，无主机-设备同步开销。
? 延伸建议：
若后续需扩展至多状态（如 [-2,-1,0,1,2]），可将三态检测泛化为：
target_values = torch.tensor([-1, 0, 1], device=x_last.device) mask = torch.stack([(x_last == v).any(dim=1) for v in target_values]).all(dim=0)
但需权衡小规模枚举 vs 可读性。当前方案已兼顾性能、简洁性与可维护性，是 PyTorch 高效数据筛选的典型范式。
