PyTorch 中是否存在干净且可扩展的 LSTM 实现? [英] Does a clean and extendable LSTM implementation exists in PyTorch?

问题描述

我想自己创建一个LSTM类，但是，我不想再次从头开始重写经典的LSTM函数.

I would like to create an LSTM class by myself, however, I don't want to rewrite the classic LSTM functions from scratch again.

在PyTorch的代码中挖掘，我只找到一个涉及至少3-4个继承类的脏实现:

Digging in the code of PyTorch, I only find a dirty implementation involving at least 3-4 classes with inheritance:

1. https://github/pytorch/pytorch/blob/98c24fae6b6400a7d1e13610b20aa05f86f77070/torch/nn/modules/rnn.py#L323
2. https://github.com/pytorch/pytorch/blob/98c24fae6b6400a7d1e13610b20aa05f86f77070/torch/nn/modules/rnn.py#L12
3. https://https://github.com/pytorch/pytorch/blob/98c24fae6b6400a7d1e13610b20aa05f86f77070/torch/nn/_functions/rnn./pytorch/pytorch/blob/98c24fae6b6400a7d1e13610b20aa05f86f77070/torch/nn/_functions/rnn.py#L297

是否存在LSTM 的clean PyTorch 实现?任何链接都会有所帮助.

Does a clean PyTorch implementation of an LSTM exist somewhere? Any links would help.

例如，我知道 LSTM 的干净实现存在于 TensorFlow 中，但我需要派生一个 PyTorch.

For example, I know that clean implementations of a LSTM exists in TensorFlow, but I would need to derive a PyTorch one.

举个明显的例子，我正在寻找的是一个像 this，但在 PyTorch 中:

For a clear example, what I'm searching for is an implementation as clean as this, but in PyTorch:

推荐答案

我找到的最好的实现在这里
https://github.com/pytorch/benchmark/blob/master/rnns/benchmarks/lstm_variants/lstm.py

The best implementation I found is here
https://github.com/pytorch/benchmark/blob/master/rnns/benchmarks/lstm_variants/lstm.py

它甚至实现了四种不同的recurrent dropout变体，非常有用！
如果你把辍学部分拿走，你会得到

It even implements four different variants of recurrent dropout, which is very useful!
If you take the dropout parts away you get

import math
import torch as th
import torch.nn as nn

class LSTM(nn.Module):

    def __init__(self, input_size, hidden_size, bias=True):
        super(LSTM, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.bias = bias
        self.i2h = nn.Linear(input_size, 4 * hidden_size, bias=bias)
        self.h2h = nn.Linear(hidden_size, 4 * hidden_size, bias=bias)
        self.reset_parameters()

    def reset_parameters(self):
        std = 1.0 / math.sqrt(self.hidden_size)
        for w in self.parameters():
            w.data.uniform_(-std, std)

    def forward(self, x, hidden):
        h, c = hidden
        h = h.view(h.size(1), -1)
        c = c.view(c.size(1), -1)
        x = x.view(x.size(1), -1)

        # Linear mappings
        preact = self.i2h(x) + self.h2h(h)

        # activations
        gates = preact[:, :3 * self.hidden_size].sigmoid()
        g_t = preact[:, 3 * self.hidden_size:].tanh()
        i_t = gates[:, :self.hidden_size]
        f_t = gates[:, self.hidden_size:2 * self.hidden_size]
        o_t = gates[:, -self.hidden_size:]

        c_t = th.mul(c, f_t) + th.mul(i_t, g_t)

        h_t = th.mul(o_t, c_t.tanh())

        h_t = h_t.view(1, h_t.size(0), -1)
        c_t = c_t.view(1, c_t.size(0), -1)
        return h_t, (h_t, c_t)

PS:存储库包含更多 LSTM 和其他 RNN 的变体:
https://github.com/pytorch/benchmark/tree/master/rnns/基准.
看看吧，也许你心目中的扩展已经有了！

PS: The repository contains many more variants of LSTM and other RNNs:
https://github.com/pytorch/benchmark/tree/master/rnns/benchmarks.
Check it out, maybe the extension you had in mind is already there!

正如评论中提到的，您可以将上面的 LSTM 单元包装起来以处理顺序输出:

As mentioned in the comments, you can wrap the LSTM cell above to process sequential output:

import math
import torch as th
import torch.nn as nn


class LSTMCell(nn.Module):

    def __init__(self, input_size, hidden_size, bias=True):
        # As before

    def reset_parameters(self):
        # As before

    def forward(self, x, hidden):

        if hidden is None:
            hidden = self._init_hidden(x)

        # Rest as before

    @staticmethod
    def _init_hidden(input_):
        h = th.zeros_like(input_.view(1, input_.size(1), -1))
        c = th.zeros_like(input_.view(1, input_.size(1), -1))
        return h, c


class LSTM(nn.Module):

    def __init__(self, input_size, hidden_size, bias=True):
        super().__init__()
        self.lstm_cell = LSTMCell(input_size, hidden_size, bias)

    def forward(self, input_, hidden=None):
        # input_ is of dimensionalty (1, time, input_size, ...)

        outputs = []
        for x in torch.unbind(input_, dim=1):
            hidden = self.lstm_cell(x, hidden)
            outputs.append(hidden[0].clone())

        return torch.stack(outputs, dim=1)

我没有测试代码，因为我正在使用 convLSTM 实现.如果有什么问题，请告诉我.

I havn't tested the code since I'm working with a convLSTM implementation. Please let me know if something is wrong.

更新:固定链接.

这篇关于PyTorch 中是否存在干净且可扩展的 LSTM 实现?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！