From a0031324174d8e6b44b633f8213fffe8b222acb2 Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Thu, 16 Jan 2020 10:47:09 -0800 Subject: [PATCH] fix lstm layer with projection save params (#17266) (#17288) Co-authored-by: Sheng Zha --- python/mxnet/gluon/rnn/rnn_layer.py | 2 +- tests/python/gpu/test_gluon_gpu.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/python/mxnet/gluon/rnn/rnn_layer.py b/python/mxnet/gluon/rnn/rnn_layer.py index 9807c5e33108..f4489b7656d6 100644 --- a/python/mxnet/gluon/rnn/rnn_layer.py +++ b/python/mxnet/gluon/rnn/rnn_layer.py @@ -126,7 +126,7 @@ def __repr__(self): def _collect_params_with_prefix(self, prefix=''): if prefix: prefix += '.' - pattern = re.compile(r'(l|r)(\d)_(i2h|h2h)_(weight|bias)\Z') + pattern = re.compile(r'(l|r)(\d)_(i2h|h2h|h2r)_(weight|bias)\Z') def convert_key(m, bidirectional): # for compatibility with old parameter format d, l, g, t = [m.group(i) for i in range(1, 5)] if bidirectional: diff --git a/tests/python/gpu/test_gluon_gpu.py b/tests/python/gpu/test_gluon_gpu.py index b938b5783415..64a8040e8802 100644 --- a/tests/python/gpu/test_gluon_gpu.py +++ b/tests/python/gpu/test_gluon_gpu.py @@ -126,6 +126,8 @@ def test_lstmp(): check_rnn_layer_forward(gluon.rnn.LSTM(10, 2, bidirectional=True, dropout=0.5, projection_size=5), mx.nd.ones((8, 3, 20)), [mx.nd.ones((4, 3, 5)), mx.nd.ones((4, 3, 10))], run_only=True, ctx=ctx) + lstm_layer.save_parameters('gpu_tmp.params') + lstm_layer.load_parameters('gpu_tmp.params') @with_seed()