MXNet Gluon のLSTMについて

>>> model = mx.gluon.rnn.LSTM(512, num_layers=1)
>>> model.initialize()
>>> input = mx.nd.random.uniform(shape=(15,10,200))
>>> h0 = mx.nd.zeros(shape=(1,10,512))
>>> c0 = mx.nd.zeros(shape=(1,10,512))
>>> out, state = model(input, [h0,c0])
>>> out.shape
(15, 10, 512)
>>> state[0].shape
(1, 10, 512)
>>> state[1].shape
(1, 10, 512)
>>> out[14]==state[0]

[[[1. 1. 1. ... 1. 1. 1.]
  [1. 1. 1. ... 1. 1. 1.]
  [1. 1. 1. ... 1. 1. 1.]
  ...
  [1. 1. 1. ... 1. 1. 1.]
  [1. 1. 1. ... 1. 1. 1.]
  [1. 1. 1. ... 1. 1. 1.]]]
<NDArray 1x10x512 @cpu(0)>

outの最後とstate[0]は当然ながら一致する。

Seq2Seq

翻訳などの文章→文章の場合にはEncoderの出力のうち「h」のみをDecoderに渡す場合と「h」「c」の両方を渡す場合が考えられる。下記参考文献（p.297）によると「h」のみを渡すことが一般的らしい。

Image Captioning

Encoderとして学習済みのResNetを使用したモデルをMXNetで書いてみた。

import mxnet as mx
from mxnet.gluon import nn, rnn, Block
from gluoncv.model_zoo import get_model

class Encoder(Block):
    def __init__(self, num_hidden):
        super(Encoder, self).__init__()

        with self.name_scope():
            self.features = get_model('ResNet50_v2', pretrained=False).features
            self.output = nn.Dense(num_hidden)
    
    def forward(self, X):
        out = self.features(X)
        out = self.output(out)
        return out


class Decoder(Block):
    def __init__(self, vocab_size, embed_size, num_hiddens):
        super(Decoder, self).__init__()

        with self.name_scope():
            self.embedding = nn.Embedding(vocab_size, embed_size)
            self.rnn = rnn.LSTM(num_hiddens, num_layers=1)
            self.dense = nn.Dense(vocab_size, flatten=False)

    def forward(self, X, features):
        X = self.embedding(X)
        X = X.swapaxes(0, 1)
        h0 = features.expand_dims(0)
        c0 = mx.nd.zeros_like(h0)
        out, state = self.rnn(X, [h0,c0])
        out = out.swapaxes(0, 1)
        out = self.dense(out)
        return out


class EncoderDecoder(Block):
    def __init__(self, vocab_size, embed_size, num_hiddens):
        super(EncoderDecoder, self).__init__()

        with self.name_scope():
            self.encoder = Encoder(num_hiddens)
            self.decoder = Decoder(vocab_size, embed_size, num_hiddens)

    def forward(self, enc_X, dec_X):
        features = self.encoder(enc_X)
        return self.decoder(dec_X, features)

あらかじめ学習済みパラメーターを保存しておく

from gluoncv.model_zoo import get_model
net = get_model('ResNet50_v2', pretrained=True)
net.features.save_parameters('features.params')

学習済み部分とそうでない部分を別々に初期化する。

import mxnet as mx
net = model.EncoderDecoder(8000, 200, 512)
net.encoder.features.load_parameters('features.params')
net.encoder.output.initialize(mx.init.Xavier())
net.decoder.initialize(mx.init.Xavier())

「features」部分のパラメータを固定する場合には以下の一文を加える。

net.encoder.features.collect_params().setattr('grad_req', 'null')

参考文献

ゼロから作るDeep Learning ? ―自然言語処理編