fix block.export

python/mxnet/gluon/block.py

@@ -1234,7 +1234,7 @@ def infer_type(self, *args):
         """Infers data type of Parameters from inputs."""
         self._infer_attrs('infer_type', 'dtype', *args)
 
-    def export(self, path, epoch=0, remove_amp_cast=True):
+    def export(self, path, epoch=0, remove_amp_cast=True, allow_extra=True):
         """Export HybridBlock to json format that can be loaded by
         `gluon.SymbolBlock.imports`, `mxnet.mod.Module` or the C++ interface.
 
@@ -1248,7 +1248,11 @@ def export(self, path, epoch=0, remove_amp_cast=True):
             will be created, where xxxx is the 4 digits epoch number.
         epoch : int
             Epoch number of saved model.
-
+        remove_amp_cast : bool, optional
+            Whether to remove the amp_cast and amp_multicast operators, before exporting the model.
+        allow_extra : bool, optional
+            Whether to save extra parameters whose names are not in the result symbol.
+            User can set allow_extra to True to load these parameters with old mxnet.mod.Module.set_params API.
         Returns
         -------
         symbol_filename : str
@@ -1271,8 +1275,13 @@ def export(self, path, epoch=0, remove_amp_cast=True):
             if name in arg_names:
                 arg_dict['arg:%s'%name] = param._reduce()
             else:
-                assert name in aux_names
-                arg_dict['aux:%s'%name] = param._reduce()
+                if name.endswith('running_mean') or name.endswith('running_var') \
+                    or name.endswith('moving_mean') or name.endswith('moving_var'):
+                    assert name in aux_names
+                    arg_dict['aux:%s'%name] = param._reduce()
+                else:
+                    if allow_extra:
+                        arg_dict['aux:%s'%name] = param._reduce()
         save_fn = _mx_npx.save if is_np_array() else ndarray.save
         params_filename = '%s-%04d.params'%(path, epoch)
         save_fn(params_filename, arg_dict)