如何高效去除Tensor中的填充值？（附示例）

给定如下PyTorch张量（填充值为-1）：

tensor([[ 9, -1, -1],
        [ 7, -1, -1],
        [ 6,  4, -1]])

要得到每行去除-1后的不等长列表[[9], [7], [6, 4]]，这里提供两种高效实现方式：

方法一：列表推导式+布尔索引（简洁直观） #

直接遍历张量每行，用布尔索引筛选非填充值后转成Python列表，代码简洁且效率足够：

import torch

input_tensor = torch.tensor([[ 9, -1, -1],
                             [ 7, -1, -1],
                             [ 6,  4, -1]])

output = [row[row != -1].tolist() for row in input_tensor]
print(output)  # 输出: [[9], [7], [6, 4]]

方法二：批量掩码处理（大规模数据更高效） #

如果处理的张量规模较大，批量处理能避免逐行遍历的开销，步骤如下：

import torch

input_tensor = torch.tensor([[ 9, -1, -1],
                             [ 7, -1, -1],
                             [ 6,  4, -1]])

# 生成掩码：标记非填充值的位置
mask = input_tensor != -1
# 统计每行的有效元素数量
valid_counts = mask.sum(dim=1)
# 提取所有有效元素
all_valid = input_tensor[mask]
# 按原行的有效数量拆分张量
split_tensors = torch.split(all_valid, valid_counts.tolist())
# 转成列表格式
output = [t.tolist() for t in split_tensors]
print(output)  # 输出: [[9], [7], [6, 4]]

适用场景说明 #

• 中小规模张量优先选方法一，代码易读且实现快；
• 大规模张量选方法二，批量操作的性能优势更明显。

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