本文介绍如何将逐列判断张量中“活跃模式”（即在最后时间步同时包含-1、0、1三种状态）的python循环逻辑，完全向量化为高性能pytorch操作，避免for循环，提升数十倍执行效率。
本文介绍如何将逐列判断张量中“活跃模式”（即在最后时间步同时包含-1、0、1三种状态）的python循环逻辑，完全向量化为高性能pytorch操作，避免for循环，提升数十倍执行效率。
在处理高维时序张量（如 torch.Size([51, 265, 23, 23])）时，常需按通道（pattern）进行条件筛选。本例中，每个 pattern 是一个 23×23 空间矩阵，其状态值域为 {-1, 0, 1}；仅当在最后一个时间步（即 tensor_sims[-1, i]）中同时出现全部三种值时，该 pattern 才被视为“活跃”，应被保留。
原始实现使用 Python 列表推导式配合 torch.unique().numel() == 3 判断，虽逻辑清晰，但因显式遍历 range(tensor_sims.shape[1])，导致严重性能瓶颈：
# ❌ 低效：含Python for循环，无法利用GPU并行
keep_indices = torch.tensor([i for i in range(tensor_sims.shape[1])
if tensor_sims[-1, i].unique().numel() == 3])torch.unique() 不支持沿指定维度（如 dim=0）批量计算每列的唯一值数量，因此无法直接向量化该逻辑。正确思路是转为布尔掩码驱动的向量化判定：
- 聚焦目标切片：提取最后一时间步 tensor_sims[-1]，形状为 [265, 23, 23]；
- 空间展平：调用 .flatten(1) 将后两个空间维压平为一维，得到 [265, 529] 张量 x_；
-
构建三元存在性掩码：对每个 pattern（行），分别判断 -1、0、1 是否至少出现一次：
mask_minus_one = (x_ == -1).any(dim=1) # shape: [265] mask_zero = (x_ == 0).any(dim=1) # shape: [265] mask_one = (x_ == 1).any(dim=1) # shape: [265]
-
逻辑合取（AND）得最终保留掩码：仅当三者均为 True 时，该 pattern 满足“三态共存”条件：
mask = mask_minus_one & mask_zero & mask_one # 或使用 .logical_and()
-
索引筛选：用布尔掩码直接切片，完成高效过滤：
tensor_sims = tensor_sims[:, mask] # 自动保持维度一致性
完整优化函数如下：
def filter_patterns(tensor_sims):
# 提取最后时间步并展平空间维度：[265, 23, 23] → [265, 529]
x_ = tensor_sims[-1].flatten(1)
# 构建三态存在性掩码（每行独立判断）
has_neg1 = (x_ == -1).any(dim=1)
has_zero = (x_ == 0).any(dim=1)
has_pos1 = (x_ == 1).any(dim=1)
# 合取：仅当三态均存在时保留
keep_mask = has_neg1 & has_zero & has_pos1
# 向量化索引筛选（支持CUDA张量）
filtered = tensor_sims[:, keep_mask]
print(f'Number of patterns retained: {filtered.shape[1]}')
return filtered✅ 优势总结：
- 零Python循环：全程基于Tensor操作，可全自动迁移至GPU加速；
- 内存友好：避免生成中间索引列表，减少CPU-GPU数据拷贝；
- 可扩展性强：若需扩展为四态（如加入2）、或改为“至少含两态”，仅需调整掩码逻辑；
- 调试直观：各掩码变量可直接 .item() 或 .nonzero() 检查，便于验证逻辑。
⚠️ 注意事项：
- 确保输入张量 tensor_sims 至少有1个时间步（tensor_sims.shape[0] >= 1），否则 tensor_sims[-1] 报错；
- 若存在 NaN 值，== 比较会返回 False，导致误判；建议前置清理：torch.nan_to_num(tensor_sims, nan=0.0)；
- 在训练循环中频繁调用时，可将 keep_mask 缓存复用，避免重复计算。
该方法体现了PyTorch向量化设计的核心思想：用广播+归约（any, all, sum）替代标量循环，以数据并行换取极致性能。
