在 PyTorch 中沿给定轴洗牌

如何解决在 PyTorch 中沿给定轴洗牌

我有一个数据集，它加载了以下维度 [batch_size,seq_len,n_features]（例如 torch.Size([16,600,130])）。

我希望能够在不改变 PyTorch 中的批量排序或特征向量排序的情况下，沿序列长度 axis=1 对这些数据进行混洗。

进一步说明：举例来说，假设我的批次大小为 3，序列长度为 3，特征数为 2。

示例： tensor([[[1,1],[2,2],[3,3]],[[4,4],[5,5],[6,6]],[[7,7],[8,8],[9,9]]]) 我希望能够通过以下方式随机洗牌：

tensor([[[3,3],[1,2]],[[6,6],[4,4]],[[8,[7,9]]])

是否有任何 PyTorch 函数可以自动为我执行此操作，或者有没有人知道实现这一点的好方法是什么？

解决方法

您可以使用 torch.randperm。

对于张量 t，您可以使用：

t[:,torch.randperm(t.shape[1]),:]

例如：

>>> t = torch.tensor([[[1,1],[2,2],[3,3]],[[4,4],[5,5],[6,6]],[[7,7],[8,8],[9,9]]])
>>> t
tensor([[[1,9]]])
>>> t[:,:]
tensor([[[2,3],[1,1]],[[5,6],[4,4]],[[8,9],[7,7]]])