PyTorch：如何基于输入长度张量创建样本权重矩阵

你可以利用PyTorch的广播机制和torch.arange实现完全矢量化的操作，无需任何for循环，效率拉满：

核心实现代码 #

import torch

# 输入长度张量
input_lengths = torch.tensor([4, 6, 8, 10])
# 从特征张量获取最大序列长度（也可以手动指定）
max_seq_len = 10  # 如果有特征张量，用features.shape[1]更灵活

# 生成时间步索引，形状为 (max_seq_len,)
time_indices = torch.arange(max_seq_len, device=input_lengths.device)
# 广播比较：将input_lengths扩展为(4,1)，和time_indices逐元素比较
sample_weights = (time_indices < input_lengths.unsqueeze(1)).float()

# 输出结果
print(sample_weights)

输出结果 #

tensor([[1., 1., 1., 1., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 1., 1., 1., 0., 0., 0., 0.],
        [1., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

关键说明 #

1. 广播机制：input_lengths.unsqueeze(1)将原形状(4,)的张量扩展为(4,1)，可以和形状(10,)的time_indices自动广播为(4,10)的维度进行逐元素比较。
2. 设备对齐：指定device=input_lengths.device可以确保生成的索引张量和输入张量在同一设备（CPU/GPU）上，避免数据迁移开销。
3. 类型转换：最后用.float()将布尔张量转为浮点型权重，也可以根据需求用.long()转为整数型。

如果你的特征张量已经存在，直接用max_seq_len = features.shape[1]即可自动适配序列长度，无需手动指定。

内容的提问来源于stack exchange，提问作者Kaan Dönbekci