如何在PyTorch中对张量列表左侧补零至最长张量长度？

在PyTorch中，使用torch.nn.utils.rnn.pad_sequence可以快速实现张量列表的右侧补零，示例代码及输出如下：

import torch

f = [
    [0,1],
    [0, 3, 4],
    [4, 3, 2, 4, 3]
]

torch.nn.utils.rnn.pad_sequence(
    [torch.tensor(part) for part in f],
    batch_first=True
)

输出结果：

tensor([[0, 1, 0, 0, 0],
        [0, 3, 4, 0, 0],
        [4, 3, 2, 4, 3]])

若要实现左侧补零，得到如下期望输出：

tensor([[0, 0, 0, 0, 1],
        [0, 0, 0, 3, 4],
        [4, 3, 2, 4, 3]])

可以通过以下两种方法实现：

方法一：翻转序列后补零再翻转 #

利用pad_sequence的右侧补零逻辑，先将每个序列翻转，补零后再翻转回来，间接实现左侧补零：

import torch

f = [
    [0,1],
    [0, 3, 4],
    [4, 3, 2, 4, 3]
]

# 转换为张量并翻转每个序列
tensors = [torch.tensor(seq) for seq in f]
reversed_tensors = [t.flip(dims=[0]) for t in tensors]

# 对翻转后的序列右侧补零
padded_reversed = torch.nn.utils.rnn.pad_sequence(reversed_tensors, batch_first=True)

# 再次翻转得到左侧补零的结果
padded_left = padded_reversed.flip(dims=[1])

print(padded_left)

方法二：手动计算padding并拼接零张量 #

直接计算每个序列需要补零的长度，手动拼接零张量和原序列，最后堆叠成批量张量：

import torch

f = [
    [0,1],
    [0, 3, 4],
    [4, 3, 2, 4, 3]
]

tensors = [torch.tensor(seq) for seq in f]
max_seq_len = max(t.shape[0] for t in tensors)
padded_tensors = []

for t in tensors:
    pad_length = max_seq_len - t.shape[0]
    # 生成左侧补零的张量并拼接
    padded_tensor = torch.cat([torch.zeros(pad_length, dtype=t.dtype), t], dim=0)
    padded_tensors.append(padded_tensor)

# 堆叠所有补零后的张量
padded_left = torch.stack(padded_tensors)

print(padded_left)

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