web-dev-qa-db-fra.com

Existe-t-il une implémentation LSTM propre et extensible dans PyTorch?

Je voudrais créer une classe LSTM par moi-même, cependant, je ne veux pas réécrire les fonctions LSTM classiques à partir de zéro.

En creusant dans le code de PyTorch, je ne trouve qu'une implémentation sale impliquant au moins 3-4 classes avec héritage: https://github.com/pytorch/pytorch/blob/98c24fae6b6400a7d1e13610b20aa05f86f77070/torch/nn/modules/rnn .py # L32https://github.com/pytorch/pytorch/blob/98c24fae6b6400a7d1e13610b20aa05f86f77070/torch/nn/modules/rnn.py#L12https: // github .com/pytorch/pytorch/blob/98c24fae6b6400a7d1e13610b20aa05f86f77070/torch/nn/_functions/rnn.py # L297

Existe-t-il une implémentation PyTorch propre d'un LSTM quelque part? Tous les liens seraient utiles. Par exemple, je sais que des implémentations propres d'un LSTM existent dans TensorFlow, mais j'aurais besoin d'en dériver un PyTorch.

Pour un exemple clair, ce que je recherche est une implémentation aussi propre que la suivante, mais dans PyTorch: https://github.com/hardmaru/supercell/blob/063b01e75e6e8af5aeb0aac5cc583948f5887dd1/supercell.py#L14

10

La meilleure implémentation que j'ai trouvée est ici
https://github.com/pytorch/benchmark/blob/master/benchmarks/lstm_variants/lstm.py

Il implémente même quatre variantes différentes de décrochage récurrent, ce qui est très utile!
Si vous enlevez les pièces de décrochage, vous obtenez

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: Le référentiel contient de nombreuses autres variantes de LSTM et d'autres RNN:
https://github.com/pytorch/benchmark/tree/master/benchmarks .
Vérifiez, peut-être que l'extension que vous aviez en tête est déjà là!

ÉDITER:
Comme mentionné dans les commentaires, vous pouvez encapsuler la cellule LSTM ci-dessus pour traiter la sortie séquentielle:

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)

Je n'ai pas testé le code depuis que je travaille avec une implémentation convLSTM. Veuillez me faire savoir si quelque chose ne va pas.

13
Richard

J'ai fait un cadre simple et général pour personnaliser les LSTM: https://github.com/daehwannam/pytorch-rnn-util

Vous pouvez implémenter des LSTM personnalisés en concevant des cellules LSTM et en les fournissant à LSTMFrame . Un exemple de LSTM personnalisé est LayerNormLSTM dans le package:

# snippet from rnn_util/seq.py
class LayerNormLSTM(LSTMFrame):
    def __init__(self, input_size, hidden_size, num_layers=1, dropout=0, r_dropout=0, bidirectional=False, layer_norm_enabled=True):
        r_dropout_layer = nn.Dropout(r_dropout)
        rnn_cells = Tuple(
            Tuple(
                LayerNormLSTMCell(
                    input_size if layer_idx == 0 else hidden_size * (2 if bidirectional else 1),
                    hidden_size,
                    dropout=r_dropout_layer,
                    layer_norm_enabled=layer_norm_enabled)
                for _ in range(2 if bidirectional else 1))
            for layer_idx in range(num_layers))

        super().__init__(rnn_cells, dropout, bidirectional)

LayerNormLSTM a les options clés du LSTM standard de PyTorch et des options supplémentaires, r_dropout et layer_norm_enabled:

# example.py
import torch
import rnn_util


bidirectional = True
num_directions = 2 if bidirectional else 1

rnn = rnn_util.LayerNormLSTM(10, 20, 2, dropout=0.3, r_dropout=0.25,
                             bidirectional=bidirectional, layer_norm_enabled=True)
# rnn = torch.nn.LSTM(10, 20, 2, bidirectional=bidirectional)

input = torch.randn(5, 3, 10)
h0 = torch.randn(2 * num_directions, 3, 20)
c0 = torch.randn(2 * num_directions, 3, 20)
output, (hn, cn) = rnn(input, (h0, c0))

print(output.size())
1
dhnam