Python 自定义Pytorch层，在每组上应用LSTM

我有一个带特征的

N×F

张量和一个带组索引的

N×1

张量。我想设计一个定制的

pytorch

层，它将

LSTM

应用于每个具有排序特征的组。我已经提到了以排序组特性为例的

LSTM

，假设它可以是任何支持可变长度输入或序列的东西。请参考下面的图片，了解问题的视觉解释

---

---

---

显而易见的方法是为每个独特的组调用

LSTM

层，但这样做效率很低。有更好的方法吗？

您当然可以并行化LSTM应用程序——问题是有效地索引特征张量。 我能想到的最好的办法（我用一些类似的方法来处理我自己的东西）是在唯一的组ID上列出一个可变长度张量的列表，然后把它们放在上面，在上面运行LSTM

代码：

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

n = 13
f = 77
n_groups = 3

xs = torch.rand(n, f)
ids = torch.randint(low=0, high=n_groups, size=(n,))


def groupbyid(xs: Tensor, ids: Tensor, batch_first: bool, 
              padding_value: int = 0) -> Tensor:
   return pad_sequence([xs[ids==idx] for idx in ids.unique()], 
                       batch_first=batch_first, 
                       padding_value=padding_value)

grouped = groupbyid(xs, ids)
print(grouped.shape)  
# torch.Size([3, 5, 77]) 

然后可以在

分组

张量上的

n_组

维度上并行应用LSTM

---

请注意您还需要检查

ids.unique（）

的内容，以将每个LSTM输出分配给其相应的组id，但这很容易编写，并且取决于您的应用程序。

在这种情况下是否需要填充？否则，您无法在可变长度序列上本机应用LSTM