fix lstm layer with projection save params (apache#17266)

frankfliu · Jan 13, 2020 · 52fc0a3 · 52fc0a3
1 parent de84f23
commit 52fc0a3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git 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
@@ -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()