冒泡 3-D 张量 PyTorch 的非空行

Question

我在尝试将以下函数转换为仅处理张量的函数时遇到了一些问题。

def valid_sequence_output(sequence_output, valid_mask):
    bs, max_len, feat_dim = sequence_output.shape
    valid_output = torch.zeros(bs, max_len, feat_dim, dtype=torch.float32)
    for i in range(bs):
        jj = -1
        for j in range(max_len):
            if valid_mask[i][j].item() == 1:
                jj += 1
                valid_output[i][jj] = sequence_output[i][j]
    return valid_output

可以按如下方式创建输入张量：

size = ((2,5,2))
sequence_output = torch.randint(0, 250, size=size)
valid_mask = torch.randint(0, 2, size=size[:2])

我的目标基本上是“冒泡” sequence_output 的非空行。例如 sequence_output 等于：

tensor([[[  0,   0],
         [ 15,  47],
         [124, 230],
         [ 0,  0],
         [  65,   31]],
        [[  0,   0],
         [  0,   0],
         [  0,   0],
         [139, 228],
         [224, 205]]])

我正在尝试获取以下张量：

tensor([[[ 15,  47],
         [124, 230],
         [ 65,  31],
         [  0,   0],
         [  0,   0]],
        [[139, 228],
         [224, 205],
         [  0,   0],
         [  0,   0],
         [  0,   0]]])

如果有人对如何执行此操作有建议，我将不胜感激:D

Answer 1

我想出了一个非常糟糕的解决方案（他可能不是最优的），方法是构建一个由零和零组成的矩阵 A，在执行矩阵乘法时将交换 X 的行。 ..

def vso(seq_out, valid_mask):
    X = torch.where(valid_mask.unsqueeze(-1) == 1, seq_out, torch.zeros_like(seq_out))
    bs, max_len, _ = X.shape
    tu = torch.unique(torch.nonzero(X)[:, :2], dim=0)
    batch_axis = tu[:, 0]
    rows_axis = tu[:, 1]
    a = torch.arange(bs).repeat(batch_axis.shape).reshape(batch_axis.shape[0], -1).T
    T = torch.cumsum(batch_axis == a, dim=1) - 1
    cols_axis = T[batch_axis, torch.arange(batch_axis.shape[0])]
    A = torch.zeros((bs, max_len, max_len))
    A[(batch_axis, cols_axis, rows_axis)] = 1
    valid_output = torch.matmul(A, X)
    return valid_output

仍在寻找更好的答案！

冒泡 3-D 张量 PyTorch 的非空行

Bubble up non null rows of a 3-D Tensor PyTorch

python

linear-algebra

pytorch

tensor