如何在PyTorch中更新3D张量第二维度首个符合条件的元素？

需求：给定一个3D张量，需要将第二维度下每个向量中首个满足条件（val==1）的元素更新为2。已知每个第二维度的向量均以0开头，且至少包含一个1。

示例输入 #

x = torch.tensor([[[   0, 2239, 6255,    1,    1],
         [   0,    1,    1,    1,    1],
         [   0,    1,    1,    1,    1]],

        [[   0, 7596, 2027, 2007,    1],
         [   0, 2030, 6255, 2752,    1],
         [   0,    1,    1,    1,    1]],

        [[   0, 2027, 2007, 2030,    1],
         [   0, 2064, 6255, 2012,    1],
         [   0, 2007, 2940,    1,    1]],

        [[   0, 2027, 2007, 2030,    1],
         [   0, 2064, 6255, 2012,    1],
         [   0, 2007, 2940, 2940,    1]]])

期望输出 #

tensor([[[   0, 2239, 6255,    2,    1],
         [   0,    2,    1,    1,    1],
         [   0,    2,    1,    1,    1]],

        [[   0, 7596, 2027, 2007,    2],
         [   0, 2030, 6255, 2752,    2],
         [   0,    2,    1,    1,    1]],

        [[   0, 2027, 2007, 2030,    2],
         [   0, 2064, 6255, 2012,    2],
         [   0, 2007, 2940,    2,    1]],

        [[   0, 2027, 2007, 2030,    2],
         [   0, 2064, 6255, 2012,    2],
         [   0, 2007, 2940, 2940,    2]]])

用户尝试的代码 #

目前已找出所有值为1的元素索引，但无法筛选出第二维度下的首个匹配项：

ids = torch.argwhere(x==1)

# TODO: code to filter ids to the first match across dimensions 2

x[ids[:, 0], ids[:, 1], ids[:, 2]] = 2

解决方案 #

方法一：利用argmax定位首个1的位置 #

argmax在布尔张量（转为int后）中会返回**第一个最大值（即第一个1）**的索引，结合高级索引即可完成更新：

import torch

# 生成标记值为1的布尔掩码
mask = x == 1
# 找到每个(batch, seq)对应的首个1在第三维度的索引
first_ones_idx = mask.int().argmax(dim=2)

# 构造第一维和第二维度的索引矩阵，与first_ones_idx匹配
batch_idx = torch.arange(x.size(0)).unsqueeze(1).repeat(1, x.size(1))
seq_idx = torch.arange(x.size(1)).unsqueeze(0).repeat(x.size(0), 1)

# 定位并更新首个1为2
x[batch_idx, seq_idx, first_ones_idx] = 2

方法二：利用cummax生成首个1的掩码 #

通过累积最大值筛选出每个向量中第一个出现的1，生成掩码后直接更新：

import torch

mask = x == 1
# 累积最大值，标记当前及之前是否出现过1
cum_mask = mask.cummax(dim=2)[0]
# 仅保留第一个1的位置：当前是1，且之前没有出现过1
first_ones_mask = mask & ~cum_mask.roll(shifts=1, dims=2)
# 修正第一个元素（向量以0开头，第一个位置不可能是1）
first_ones_mask[:, :, 0] = False

# 用掩码更新张量
x[first_ones_mask] = 2

两种方法都能高效完成需求，其中方法一代码更简洁，方法二逻辑更直观，可根据实际场景选择。

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