PyTorch中如何沿seq_len维度（axis=1）打乱张量且保留批次和特征顺序

PyTorch没有直接实现该需求的内置函数，你可以通过随机索引重排的方式快速实现，代码可直接跑在CPU/GPU上，没有额外效率损耗。

两种常见实现场景 #

场景1：每个batch的序列使用独立打乱规则（符合你给出的示例需求） #

适用于需要每个样本的序列打乱顺序完全独立的场景，实现代码如下：

import torch

# 构造示例输入
x = torch.tensor([[[1,1],[2,2],[3,3]],
                  [[4,4],[5,5],[6,6]],
                  [[7,7],[8,8],[9,9]]])
batch_size, seq_len, n_features = x.shape

# 生成每个batch独立的随机索引
perm = torch.argsort(torch.rand(batch_size, seq_len), dim=1)
# 扩展索引维度匹配输入张量
perm = perm.unsqueeze(-1).expand(-1, -1, n_features)
# 沿序列维度（dim=1）重排
x_shuffled = torch.gather(x, dim=1, index=perm)

场景2：所有batch使用相同的序列打乱规则 #

适用于要求同一批所有样本的序列打乱规则一致的场景，实现更简单：

seq_len = x.shape[1]
# 生成统一的随机索引
perm = torch.randperm(seq_len)
# 直接索引重排
x_shuffled = x[:, perm, :]

结果验证 #

以你给出的示例输入为例，场景1的输出会和你给出的期望结果结构一致，每个batch的序列顺序随机打乱，批次顺序、特征向量内部顺序完全不变。

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