如何转换 torch.nn.utils.rnn.pack_sequence 的 2D 和索引张量

Question

我有以下形式的序列集合：

sequences = torch.tensor([[2,1],[5,6],[3,0])
indexes = torch.tensor([1,0,1])

即序列0仅由[5,6]组成，序列1由[2,1] , [3,0]组成。数学上 sequence[i] = { sequences[j] such that i = indexes[j] }

我需要将这些序列输入 LSTM. Since these are variable-length sequences, pytorch documentation states to use something like torch.nn.utils.rnn.pack_sequence。

遗憾的是，此方法及其类似方法需要一个张量列表作为输入，其中每个张量都是一个 L x *，其中 L 是单个序列的长度。

如何构建可以输入 pytorch LSTM 的东西？

P.s。在整个代码中，我使用 scatter 和 gather 功能来处理这些张量，但我找不到使用它们来实现此目标的方法。

Answer 1

首先，你需要分离你的序列。 Pack_sequence 接受张量列表，每个张量的形状为 L x *。所有序列的其他维度必须始终相同，但 L 或序列长度可以变化。比如你的序列0和1可以打包为：

sequences = [torch.tensor([[5,6]]), torch.tensor([[2,1],[3,0]])]
packed_seq = torch.nn.utils.rnn.pack_sequence(sequences, enforce_sorted=False)

这里，在sequences中，sequences[0]的shape是(1,2)，而sequences[1]的shape是(2,2)。第一维表示它们的长度，分别为1和2。

您可以通过以下方式分隔序列：

sequences = torch.tensor([[2,1],[5,6],[3,0]])
indexes = torch.tensor([1,0,1])
num_seq = np.unique(indexes)
sequences = [sequences[indexes==seq_id] for seq_id in num_seq]

这将创建 sequences=[torch.tensor([[5,6]]), torch.tensor([[2,1],[3,0]])]。

Answer 2

我找到了另一种更有效的分离序列的方法：

sequences = torch.tensor([[2,1],[5,6],[3,0]])
indexes = torch.tensor([1,0,1])

sorted_src = src[indexes.argsort()]
indexes_count = torch.unique(indexes, return_counts=True)[1]

splitted = torch.split(sorted_src, indexes_count.tolist(), dim=0)

这种方法比@Mercury 提出的方法快将近 3 倍。

Measured using timeit module with sequences being a (5000,256) tensor and indexes being (1500)

如何转换 torch.nn.utils.rnn.pack_sequence 的 2D 和索引张量

How to transform a 2D and index tensors for torch.nn.utils.rnn.pack_sequence

python

numpy

machine-learning

lstm

pytorch