如何在PyTorch中高效过滤张量非零数据（保留批量维度）

给定形状为[batch_size, size]的张量，需要在每个批量样本中过滤掉零元素、保留批量维度，这里提供两种高效实现方案：

方案1：列表推导式（简洁直观） #

直接遍历每个批量样本，用布尔索引筛选非零元素，代码简单易懂，适合中小规模数据：

import torch

# 定义输入张量
X = torch.Tensor([[ 0, 10, 10,  0,  5,  5,  0,  0,  9,  0,  0, 20,  0,  7,  7,  0,  0,  1,
          1,  0,  0, 17, 17,  0,  0, 13, 13,  0,  0,  4,  4,  0, 18, 18,  0, 12,
         12,  0,  0],
 [ 0,  6,  6,  0,  0,  5,  0, 15,  0,  0, 11,  0,  8,  0,  0, 10,  0,  0,
          4,  0, 20, 20,  0,  0,  0, 16, 16,  0,  0, 17, 17,  0,  0, 18,  0,  0,
         19, 19,  0]])

# 过滤每个样本的非零元素
X_filtered = [x[x != 0] for x in X]

# 若需要统一形状的张量（用0填充），可使用pad_sequence
from torch.nn.utils.rnn import pad_sequence
X_filtered_padded = pad_sequence(X_filtered, batch_first=True, padding_value=0)

方案2：向量化操作（高效大规模） #

利用PyTorch的向量化API避免显式循环，底层优化更充分，适合处理大张量：

# 生成非零元素掩码
mask = X != 0
# 计算每个样本的非零元素数量
non_zero_counts = mask.sum(dim=1)
# 提取所有非零元素并按样本拆分
all_non_zeros = X[mask]
X_filtered = torch.split(all_non_zeros, non_zero_counts.tolist())

说明 #

• 两种方案最终得到的X_filtered是一个张量列表，每个元素对应原批量中一个样本的非零元素集合，与你给出的示例结构一致。
• 如果后续需要输入到深度学习模型，推荐用pad_sequence将列表转为统一形状的张量，方便批量计算。

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