自定义 Pytorch 层以在每个组上应用 LSTM答案

【问题标题】：Custom Pytorch layer to apply LSTM on each group自定义 Pytorch 层以在每个组上应用 LSTM
【发布时间】：2021-07-31 22:15:46
【问题描述】：

我有一个带有特征的 N × F 张量和一个带有组索引的 N × 1 张量。我想设计一个自定义的pytorch 层，它将LSTM 应用于具有排序功能的每个组。我曾提到LSTM 以排序组特征为例，假设它可以是任何支持可变长度输入或序列的东西。有关问题的直观解释，请参阅下图。

显而易见的方法是为每个唯一组调用LSTM 层，但这会效率低下。有没有更好的办法？

【问题讨论】：

标签： python neural-network pytorch lstm

【解决方案1】：

您当然可以并行化 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])

然后您可以在grouped 张量上的n_groups 维度上并行应用 LSTM。

请注意，您还需要检查 ids.unique() 的内容以将每个 LSTM 输出分配给其相应的组 id，但这很容易编写并且取决于您的应用程序。

【讨论】：

在这种情况下是否需要填充？
否则你不能将 LSTM 应用到可变长度序列上