torch 填充补齐

目录

行填充补齐

1. 填充长度（Padding）

2. 掩码（Masking）

3. 排序优化（可选）

---

行填充补齐

import torch
from torch.nn.utils.rnn import pad_sequence

# 原始序列（每个序列是二维张量，行数不同）
batch_data = [torch.tensor([[1, 2, 3]])  # 1行
    # torch.tensor([[4, 5, 6], [7, 8, 9]]),  # 2行
    # torch.tensor([[10, 11, 12]])  # 1行
]

# 目标：每个序列填充到5行，填充内容为 [0, 0, 21]

# 1. 创建填充模板 (每行填充内容)
padding_row = torch.tensor([[0, 0, 21]], dtype=torch.int32)

# 2. 对每个序列进行填充
padded_batch = []
for seq in batch_data:
    current_length = seq.size(0)
    padding_needed = 5 - current_length

    # 创建填充部分
    if padding_needed > 0:
        padding