MXNET CNN+LSTM save/serialize 至 json

MXNET CNN+LSTM save/serialize to json

我发现很难弄清楚如何正确定义 mxnet 网络,以便我可以 serialize/convert 这个模型到 json 文件。

管道由一个CNN + biLSTM + CTC组成。

我现在必须使用 HybridBlock 和 hybridize() 但我似乎无法使它工作,或者即使它可能,或者是否有任何其他解决方法。

我确信它缺乏我的知识,不知道是否有人可以提供帮助。

这里是python中的网络定义:

NUM_HIDDEN = 200
NUM_CLASSES = 13550
NUM_LSTM_LAYER = 1
p_dropout = 0.5
SEQ_LEN = 32

def get_featurizer():
    featurizer = gluon.nn.HybridSequential()
    # conv layer
    featurizer.add(gluon.nn.Conv2D(kernel_size=(3,3), padding=(1,1), channels=32, activation="relu"))
    featurizer.add(gluon.nn.BatchNorm())

    ....
    featurizer.hybridize()
    return featurizer

class EncoderLayer(gluon.Block):
    def __init__(self, **kwargs):
        super(EncoderLayer, self).__init__(**kwargs)
        with self.name_scope():
            self.lstm = mx.gluon.rnn.LSTM(NUM_HIDDEN, NUM_LSTM_LAYER, bidirectional=True)
    def forward(self, x):
        x = x.transpose((0,3,1,2))
        x = x.flatten()
        x = x.split(num_outputs=SEQ_LEN, axis = 1) # (SEQ_LEN, N, CHANNELS)
        x = nd.concat(*[elem.expand_dims(axis=0) for elem in x], dim=0)
        x = self.lstm(x)
        x = x.transpose((1, 0, 2)) # (N, SEQ_LEN, HIDDEN_UNITS)
        return x

def get_encoder():
    encoder = gluon.nn.Sequential()
    encoder.add(EncoderLayer())
    encoder.add(gluon.nn.Dropout(p_dropout))
    return encoder

def get_decoder():
    decoder = mx.gluon.nn.Dense(units=ALPHABET_SIZE, flatten=False)
    decoder.hybridize()
    return decoder

def get_net():
    net = gluon.nn.Sequential()
    with net.name_scope():
        net.add(get_featurizer())
        net.add(get_encoder())
        net.add(get_decoder())
    return net

如有任何帮助,我们将不胜感激。 非常感谢。

将 Gluon 中的模型导出到 json 的要求很少:

  1. 它需要是可混合的,这意味着每个子块也应该是可混合的,并且模型在两种模式下都可以工作

  2. 应初始化所有参数。由于 Gluon 使用延迟参数初始化,这意味着您应该至少进行一次正向传递才能保存模型。

我对您的代码进行了一些修复,并在需要时引入了新常量。最重要的变化是:

  1. 如果可以避免,请不要使用拆分,因为它是 returns NDArray 列表。使用 reshape,它也可以与 Symbol 完美配合。

  2. 从 1.3.0 版的 MXNet 开始,LSTM 也是可混合的,因此您可以将其包装在 HybridBlock 而不仅仅是 Block 中。

  3. 使用混合序列。

这是调整后的代码,底部有一个示例,说明如何保存模型以及如何加载模型。您可以在 this tutorial.

中找到更多信息
import mxnet as mx
from mxnet import gluon
from mxnet import nd

BATCH_SIZE = 1
CHANNELS = 100
ALPHABET_SIZE = 1000
NUM_HIDDEN = 200
NUM_CLASSES = 13550
NUM_LSTM_LAYER = 1
p_dropout = 0.5
SEQ_LEN = 32
HEIGHT = 100
WIDTH = 100


def get_featurizer():
    featurizer = gluon.nn.HybridSequential()
    featurizer.add(
        gluon.nn.Conv2D(kernel_size=(3, 3), padding=(1, 1), channels=32, activation="relu"))
    featurizer.add(gluon.nn.BatchNorm())

    return featurizer


class EncoderLayer(gluon.HybridBlock):
    def __init__(self, **kwargs):
        super(EncoderLayer, self).__init__(**kwargs)

        with self.name_scope():
            self.lstm = mx.gluon.rnn.LSTM(NUM_HIDDEN, NUM_LSTM_LAYER, bidirectional=True)

    def hybrid_forward(self, F, x):
        x = x.transpose((0, 3, 1, 2))
        x = x.flatten()
        x = x.reshape(shape=(SEQ_LEN, -1, CHANNELS)) #x.split(num_outputs=SEQ_LEN, axis=1)  # (SEQ_LEN, N, CHANNELS)
        x = self.lstm(x)
        x = x.transpose((1, 0, 2))  # (N, SEQ_LEN, HIDDEN_UNITS)
        return x


def get_encoder():
    encoder = gluon.nn.HybridSequential()
    encoder.add(EncoderLayer())
    encoder.add(gluon.nn.Dropout(p_dropout))
    return encoder


def get_decoder():
    decoder = mx.gluon.nn.Dense(units=ALPHABET_SIZE, flatten=False)
    return decoder


def get_net():
    net = gluon.nn.HybridSequential()

    with net.name_scope():
        net.add(get_featurizer())
        net.add(get_encoder())
        net.add(get_decoder())

    return net


if __name__ == '__main__':
    net = get_net()
    net.initialize()
    net.hybridize()

    fake_data = mx.random.uniform(shape=(BATCH_SIZE, HEIGHT, WIDTH, CHANNELS))
    out = net(fake_data)

    net.export("mymodel")

    deserialized_net = gluon.nn.SymbolBlock.imports("mymodel-symbol.json", ['data'],
                                                    "mymodel-0000.params", ctx=mx.cpu())

    out2 = deserialized_net(fake_data)
    # just to check that we get the same results
    assert (out - out2).sum().asscalar() == 0