自分でLSTMクラスを作成したいのですが、従来のLSTM関数を最初から書き直したくありません。
PyTorchのコードを掘り下げて、継承のある3〜4個のクラスが関係するダーティな実装のみを見つけました: /rnn.py#L323https://github.com/pytorch/pytorch/blob/98c24fae6b6400a7d1e13610b20aa05f86f77070/torch/nn/modules/rnn.py#L12https://github.com/pytorch/pytorch/blob/98c24fae6b6400a7d1e13610b20aa05f86f77070/torch/nn/_functions/rnn.py#L297
LSTMのクリーンなPyTorch実装はどこかに存在しますか?どんなリンクも役に立ちます。たとえば、LSTMのクリーンな実装がTensorFlowに存在することは知っていますが、PyTorchを派生させる必要があります。
明確な例として、私が探しているのは次のようにクリーンな実装ですが、PyTorchの場合: https://github.com/hardmaru/supercell/blob/063b01e75e6e8af5aeb0aac5cc583948f5887dd1/supercell.py# L143
私が見つけた最良の実装はここにあります
https://github.com/pytorch/benchmark/blob/master/benchmarks/lstm_variants/lstm.py
再発ドロップアウトの4つの異なるバリアントも実装しているため、非常に便利です。
ドロップアウトパーツを取り除くと、
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/benchmarks 。
確認してください。念頭に置いていた拡張機能はすでにそこにあるかもしれません。
編集:
コメントで述べたように、上記のLSTMセルをラップして順次出力を処理できます。
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実装で作業しているので、コードをテストしていません。何か問題があれば教えてください。
LSTMをカスタマイズするためのシンプルで一般的なフレームを作成しました: https://github.com/daehwannam/pytorch-rnn-util
LSTMセルを設計し、それらをLSTMFrameに提供することにより、カスタムLSTMを実装できます。カスタムLSTMの例は、パッケージ内のLayerNormLSTMです。
# 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には、PyTorchの標準LSTMの主要なオプションと、r_dropoutと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())