如何用PyTorch函数无循环批量执行torch.where(array>0)操作？

你可以利用PyTorch的向量化操作实现无循环的批量处理，核心是先提取全局非零位置信息，再按每个样本的非零元素数量拆分结果，具体实现如下：

import torch

def batch_node_indices_no_loop(states_batch):
    # 生成大于0的掩码矩阵
    mask = states_batch > 0
    # 获取所有非零元素的批次索引和样本内索引
    batch_idx, elem_idx = torch.where(mask)
    # 统计每个样本的非零元素数量
    counts = mask.sum(dim=1)
    # 按样本拆分索引张量
    split_indices = torch.split(elem_idx, counts.tolist())
    # 转换为numpy数组，和原函数输出格式对齐
    return [idx.detach().cpu().numpy() for idx in split_indices]

代码说明： #

• 第一步生成的mask和输入states_batch形状一致（假设输入为(batch_size, num_elements)的2D张量），标记每个位置是否满足大于0的条件
• 第二步用torch.where(mask)直接提取所有符合条件的位置，batch_idx对应非零元素所属的样本序号，elem_idx对应该元素在样本内的索引
• 第三步统计每个样本的非零元素数量，作为后续拆分的长度依据
• 第四步通过torch.split将全局索引拆分为对应每个样本的索引张量
• 最后一步转换为numpy数组，和原函数的输出结果完全匹配

验证示例： #

# 测试输入
states_batch = torch.tensor([
    [0, 2, 0, 5],
    [3, 0, 1, 0],
    [0, 0, 0, 0]
])

# 原函数输出
print(batch_node_indices(states_batch))
# 输出：[array([1, 3]), array([0, 2]), array([], dtype=int64)]

# 新函数输出
print(batch_node_indices_no_loop(states_batch))
# 输出：[array([1, 3]), array([0, 2]), array([], dtype=int64)]

如果不需要转换为numpy格式，直接去掉最后一步的detach().cpu().numpy()即可返回张量列表。

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