Repository: dmlc/gluon-nlp

Issue:
apply_weight_drop required changes to integrate mxnet 1.3
our existing implementation on apply_weight_drop relied on a non-public attribute in gluon RNN blocks that were removed apache/incubator-mxnet#11482 when making the RNN blocks hybrid, and it does not have test, which will prevent us from integrating with mxnet 1.3. We need:
- [ ] fix apply_weight_drop and upgrade the mxnet dependency to allow 1.3 integration.
- [ ] add test for apply_weight_drop
- [ ] bump up minor version to 0.4 which relies on 1.3.


Discussion:

Resolving patch:
diff --git a/gluonnlp/model/utils.py b/gluonnlp/model/utils.py
--- a/gluonnlp/model/utils.py
+++ b/gluonnlp/model/utils.py
@@ -19,71 +19,128 @@
 """Building blocks and utility for models."""
 __all__ = ['apply_weight_drop']
 
+import collections
 import functools
+import re
+import warnings
 
-from mxnet.gluon import Block, HybridBlock, contrib, rnn
+from mxnet.gluon import Block, contrib, rnn
 from .parameter import WeightDropParameter
 
-def apply_weight_drop(block, local_param_name, rate, axes=(),
+# pylint: disable=too-many-nested-blocks
+def apply_weight_drop(block, local_param_regex, rate, axes=(),
                       weight_dropout_mode='training'):
-    """Apply weight drop to the parameter of a block.
+    r"""Apply weight drop to the parameter of a block.
 
     Parameters
     ----------
     block : Block or HybridBlock
         The block whose parameter is to be applied weight-drop.
-    local_param_name : str
-        The parameter name used on the block. such as 'weight'.
+    local_param_regex : str
+        The regex for parameter names used in the self.params.get(), such as 'weight'.
     rate : float
         Fraction of the input units to drop. Must be a number between 0 and 1.
     axes : tuple of int, default ()
         The axes on which dropout mask is shared. If empty, regular dropout is applied.
     weight_drop_mode : {'training', 'always'}, default 'training'
         Whether the weight dropout should be applied only at training time, or always be applied.
+
+    Examples
+    --------
+    >>> import mxnet as mx
+    >>> from mxnet import gluon
+    >>> import gluonnlp as nlp
+    >>> net = gluon.rnn.LSTM(10, num_layers=2, bidirectional=True)
+    >>> nlp.model.apply_weight_drop(net, r'.*h2h_weight\Z', 0.5)
+    >>> net.collect_params()
+    lstm0_ (
+      Parameter lstm0_l0_i2h_weight (shape=(40, 0), dtype=<class 'numpy.float32'>)
+      WeightDropParameter lstm0_l0_h2h_weight
+        (shape=(40, 10), dtype=<class 'numpy.float32'>, rate=0.5, mode=training)
+      Parameter lstm0_l0_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
+      Parameter lstm0_l0_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
+      Parameter lstm0_r0_i2h_weight (shape=(40, 0), dtype=<class 'numpy.float32'>)
+      WeightDropParameter lstm0_r0_h2h_weight
+        (shape=(40, 10), dtype=<class 'numpy.float32'>, rate=0.5, mode=training)
+      Parameter lstm0_r0_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
+      Parameter lstm0_r0_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
+      Parameter lstm0_l1_i2h_weight (shape=(40, 20), dtype=<class 'numpy.float32'>)
+      WeightDropParameter lstm0_l1_h2h_weight
+        (shape=(40, 10), dtype=<class 'numpy.float32'>, rate=0.5, mode=training)
+      Parameter lstm0_l1_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
+      Parameter lstm0_l1_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
+      Parameter lstm0_r1_i2h_weight (shape=(40, 20), dtype=<class 'numpy.float32'>)
+      WeightDropParameter lstm0_r1_h2h_weight
+        (shape=(40, 10), dtype=<class 'numpy.float32'>, rate=0.5, mode=training)
+      Parameter lstm0_r1_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
+      Parameter lstm0_r1_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
+    )
+    >>> net.initialize()
+    >>> with mx.autograd.train_mode():
+    ...     print(net(mx.nd.ones((3, 4, 5))).max())
+    [0.00488924]
+    <NDArray 1 @cpu(0)>
+    >>> with mx.autograd.train_mode():
+    ...     print(net(mx.nd.ones((3, 4, 5))).max())
+    [0.00475577]
+    <NDArray 1 @cpu(0)>
     """
     if not rate:
         return
 
-    params = block.collect_params('.*{}'.format(local_param_name))
-    for full_param_name, param in params.items():
+    existing_params = _find_params(block, local_param_regex)
+    for (local_param_name, param), \
+        (ref_params_list, ref_reg_params_list) in existing_params.items():
         dropped_param = WeightDropParameter(param, rate, weight_dropout_mode, axes)
-        param_dicts, reg_param_dicts = _find_param(block, full_param_name, local_param_name)
-        for param_dict in param_dicts:
-            param_dict[full_param_name] = dropped_param
-        for reg_param_dict in reg_param_dicts:
-            reg_param_dict[local_param_name] = dropped_param
-        local_attr = getattr(block, local_param_name)
-        if local_attr == param:
-            super(Block, block).__setattr__(local_param_name, dropped_param)
-        else:
-            if isinstance(local_attr, (list, tuple)):
-                if isinstance(local_attr, tuple):
-                    local_attr = list(local_attr)
-                for i, v in enumerate(local_attr):
-                    if v == param:
-                        local_attr[i] = dropped_param
-            elif isinstance(local_attr, dict):
-                for k, v in local_attr:
-                    if v == param:
-                        local_attr[k] = dropped_param
-            else:
-                continue
-            super(Block, block).__setattr__(local_param_name, local_attr)
-
-
-def _find_param(block, full_param_name, local_param_name):
-    param_dict_results = []
-    reg_dict_results = []
-    params = block.params
-
-    if full_param_name in block.params._params:
-        if isinstance(block, HybridBlock) and local_param_name in block._reg_params:
-            reg_dict_results.append(block._reg_params)
+        for ref_params in ref_params_list:
+            ref_params[param.name] = dropped_param
+        for ref_reg_params in ref_reg_params_list:
+            ref_reg_params[local_param_name] = dropped_param
+            if hasattr(block, local_param_name):
+                local_attr = getattr(block, local_param_name)
+                if local_attr == param:
+                    local_attr = dropped_param
+                elif isinstance(local_attr, (list, tuple)):
+                    if isinstance(local_attr, tuple):
+                        local_attr = list(local_attr)
+                    for i, v in enumerate(local_attr):
+                        if v == param:
+                            local_attr[i] = dropped_param
+                elif isinstance(local_attr, dict):
+                    for k, v in local_attr:
+                        if v == param:
+                            local_attr[k] = dropped_param
+                else:
+                    continue
+                if local_attr:
+                    super(Block, block).__setattr__(local_param_name, local_attr)
+
+# pylint: enable=too-many-nested-blocks
+def _find_params(block, local_param_regex):
+    # return {(local_param_name, parameter): (referenced_params_list,
+    #                                         referenced_reg_params_list)}
+
+    results = collections.defaultdict(lambda: ([], []))
+    pattern = re.compile(local_param_regex)
+    local_param_names = ((local_param_name, p) for local_param_name, p in block._reg_params.items()
+                         if pattern.match(local_param_name))
+
+    for local_param_name, p in local_param_names:
+        ref_params_list, ref_reg_params_list = results[(local_param_name, p)]
+        ref_reg_params_list.append(block._reg_params)
+
+        params = block._params
         while params:
-            if full_param_name in params._params:
-                param_dict_results.append(params._params)
+            if p.name in params._params:
+                ref_params_list.append(params._params)
             if params._shared:
                 params = params._shared
+                warnings.warning('When applying weight drop, target parameter {} was found '
+                                 'in a shared parameter dict. The parameter attribute of the '
+                                 'original block on which the shared parameter dict was attached '
+                                 'will not be updated with WeightDropParameter. If necessary, '
+                                 'please update the attribute manually. The likely name of the '
+                                 'attribute is ".{}"'.format(p.name, local_param_name))
             else:
                 break
 
@@ -93,11 +150,13 @@ def _find_param(block, full_param_name, local_param_name):
         elif isinstance(block._children, dict):
             children = block._children.values()
         for c in children:
-            pd, rd = _find_param(c, full_param_name, local_param_name)
-            param_dict_results.extend(pd)
-            reg_dict_results.extend(rd)
+            child_results = _find_params(c, local_param_regex)
+            for (child_p_name, child_p), (child_pd_list, child_rd_list) in child_results.items():
+                pd_list, rd_list = results[(child_p_name, child_p)]
+                pd_list.extend(child_pd_list)
+                rd_list.extend(child_rd_list)
 
-    return param_dict_results, reg_dict_results
+    return results
 
 def _get_rnn_cell(mode, num_layers, input_size, hidden_size,
                   dropout, weight_dropout,
@@ -145,6 +204,6 @@ def _get_rnn_layer(mode, num_layers, input_size, hidden_size, dropout, weight_dr
                       input_size=input_size)
 
     if weight_dropout:
-        apply_weight_drop(block, 'h2h_weight', rate=weight_dropout)
+        apply_weight_drop(block, '.*h2h_weight', rate=weight_dropout)
 
     return block


Test patch:
diff --git a/tests/unittest/test_models.py b/tests/unittest/test_models.py
--- a/tests/unittest/test_models.py
+++ b/tests/unittest/test_models.py
@@ -22,10 +22,8 @@
 import sys
 
 import mxnet as mx
+from mxnet import gluon
 import gluonnlp as nlp
-from gluonnlp.model import get_model as get_text_model
-from gluonnlp.model.train import get_cache_model
-from gluonnlp.model.train import CacheCell
 
 def eprint(*args, **kwargs):
     print(*args, file=sys.stderr, **kwargs)
@@ -42,8 +40,8 @@ def _test_pretrained_big_text_models():
     for model_name in text_models:
         eprint('testing forward for %s' % model_name)
         pretrained_dataset = pretrained_to_test.get(model_name)
-        model, _ = get_text_model(model_name, dataset_name=pretrained_dataset,
-                                  pretrained=True, root='tests/data/model/')
+        model, _ = nlp.model.get_model(model_name, dataset_name=pretrained_dataset,
+                                       pretrained=True, root='tests/data/model/')
 
         print(model)
         batch_size = 10
@@ -60,7 +58,7 @@ def test_big_text_models():
 
     for model_name in text_models:
         eprint('testing forward for %s' % model_name)
-        model, _ = get_text_model(model_name, vocab=vocab, root='tests/data/model/')
+        model, _ = nlp.model.get_model(model_name, vocab=vocab, root='tests/data/model/')
 
         print(model)
         model.collect_params().initialize()
@@ -83,8 +81,9 @@ def test_text_models():
     for model_name in text_models:
         eprint('testing forward for %s' % model_name)
         pretrained_dataset = pretrained_to_test.get(model_name)
-        model, _ = get_text_model(model_name, vocab=vocab, dataset_name=pretrained_dataset,
-                                  pretrained=pretrained_dataset is not None, root='tests/data/model/')
+        model, _ = nlp.model.get_model(model_name, vocab=vocab, dataset_name=pretrained_dataset,
+                                       pretrained=pretrained_dataset is not None,
+                                       root='tests/data/model/')
 
         print(model)
         if not pretrained_dataset:
@@ -98,8 +97,8 @@ def test_cache_models():
     datasets = ['wikitext-2']
     for name in cache_language_models:
         for dataset_name in datasets:
-            cache_cell = get_cache_model(name, dataset_name, window=1, theta=0.6,
-                                         lambdas=0.2, root='tests/data/model/')
+            cache_cell = nlp.model.train.get_cache_model(name, dataset_name, window=1, theta=0.6,
+                                                         lambdas=0.2, root='tests/data/model/')
             outs, word_history, cache_history, hidden = \
                 cache_cell(mx.nd.arange(10).reshape(10, 1), mx.nd.arange(10).reshape(10, 1), None, None)
             print(cache_cell)
@@ -120,15 +119,15 @@ def test_get_cache_model_noncache_models():
     datasets = ['wikitext-2']
     for name in language_models_params.keys():
         for dataset_name in datasets:
-            _, vocab = get_text_model(name=name, dataset_name=dataset_name, pretrained=True)
+            _, vocab = nlp.model.get_model(name=name, dataset_name=dataset_name, pretrained=True)
             ntokens = len(vocab)
 
-            cache_cell_0 = get_cache_model(name, dataset_name, window=1, theta=0.6,
-                                           lambdas=0.2, root='tests/data/model/')
+            cache_cell_0 = nlp.model.train.get_cache_model(name, dataset_name, window=1, theta=0.6,
+                                                           lambdas=0.2, root='tests/data/model/')
             print(cache_cell_0)
 
-            model, _ = get_text_model(name=name, dataset_name=dataset_name, pretrained=True)
-            cache_cell_1 = CacheCell(model, ntokens, window=1, theta=0.6, lambdas=0.2)
+            model, _ = nlp.model.get_model(name=name, dataset_name=dataset_name, pretrained=True)
+            cache_cell_1 = nlp.model.train.CacheCell(model, ntokens, window=1, theta=0.6, lambdas=0.2)
             cache_cell_1.load_params('tests/data/model/' + language_models_params.get(name))
             print(cache_cell_1)
 
@@ -149,8 +148,8 @@ def test_save_load_cache_models():
     datasets = ['wikitext-2']
     for name in cache_language_models:
         for dataset_name in datasets:
-            cache_cell = get_cache_model(name, dataset_name, window=1, theta=0.6,
-                                           lambdas=0.2, root='tests/data/model/')
+            cache_cell = nlp.model.train.get_cache_model(name, dataset_name, window=1, theta=0.6,
+                                                         lambdas=0.2, root='tests/data/model/')
             print(cache_cell)
             cache_cell.save_params('tests/data/model/' + name + '-' + dataset_name + '.params')
             cache_cell.load_params('tests/data/model/' + name + '-' + dataset_name + '.params')
@@ -207,3 +206,74 @@ def test_big_rnn_model_share_params():
     logits, hidden = eval_model(x, hidden)
     assert logits.shape == (seq_len, batch_size, vocab_size)
     mx.nd.waitall()
+
+def test_weight_drop():
+    class RefBiLSTM(gluon.Block):
+        def __init__(self, size, **kwargs):
+            super(RefBiLSTM, self).__init__(**kwargs)
+            with self.name_scope():
+                self._lstm_fwd = gluon.rnn.LSTM(size, bidirectional=False, prefix='l0')
+                self._lstm_bwd = gluon.rnn.LSTM(size, bidirectional=False, prefix='r0')
+
+        def forward(self, inpt):
+            fwd = self._lstm_fwd(inpt)
+            bwd_inpt = mx.nd.flip(inpt, 0)
+            bwd = self._lstm_bwd(bwd_inpt)
+            bwd = mx.nd.flip(bwd, 0)
+            return mx.nd.concat(fwd, bwd, dim=2)
+    net1 = RefBiLSTM(10)
+    shared_net1 = RefBiLSTM(10, params=net1.collect_params())
+
+    net2 = gluon.rnn.LSTM(10)
+    shared_net2 = gluon.rnn.LSTM(10, params=net2.collect_params())
+
+    net3 = gluon.nn.HybridSequential()
+    net3.add(gluon.rnn.LSTM(10))
+    shared_net3 = gluon.nn.HybridSequential(params=net3.collect_params())
+    shared_net3.add(gluon.rnn.LSTM(10, params=net3[0].collect_params()))
+
+    x = mx.nd.ones((3, 4, 5))
+    nets = [(net1, shared_net1),
+            (net2, shared_net2),
+            (net3, shared_net3)]
+    for net, shared_net in nets:
+        net.initialize('ones')
+        mx.test_utils.assert_almost_equal(net(x).asnumpy(),
+                                          shared_net(x).asnumpy())
+        with mx.autograd.train_mode():
+            mx.test_utils.assert_almost_equal(net(x).asnumpy(),
+                                              shared_net(x).asnumpy())
+
+        grads = {}
+        with mx.autograd.record():
+            y = net(x)
+        y.backward()
+        for name, param in net.collect_params().items():
+            grads[name] = param.grad().copy()
+        with mx.autograd.record():
+            y = shared_net(x)
+        y.backward()
+        for name, param in shared_net.collect_params().items():
+            mx.test_utils.assert_almost_equal(grads[name].asnumpy(), param.grad().asnumpy())
+
+        drop_rate = 0.5
+        nlp.model.utils.apply_weight_drop(net, '.*h2h_weight', drop_rate)
+        net.initialize('ones')
+
+        mx.test_utils.assert_almost_equal(net(x).asnumpy(),
+                                          shared_net(x).asnumpy())
+        with mx.autograd.train_mode():
+            assert not mx.test_utils.almost_equal(net(x).asnumpy(),
+                                                  shared_net(x).asnumpy())
+
+        grads = {}
+        with mx.autograd.record():
+            y = net(x)
+        y.backward()
+        for name, param in net.collect_params().items():
+            grads[name] = param.grad().copy()
+        with mx.autograd.record():
+            y = shared_net(x)
+        y.backward()
+        for name, param in shared_net.collect_params().items():
+            assert not mx.test_utils.almost_equal(grads[name].asnumpy(), param.grad().asnumpy())