fix block.export (apache#17970)

* fix block.export ```net.hybridize``` may optimize out some ops. These ops are alive in nn.Block(also nn.HybridBlock), but its names are not contained in symbol's ```arg_names``` list. So ignore these ops except that their name are end with 'running_mean' or 'running_var'. * Update block.py let user can save their extra param. * add allow_extra add allow_extra to let user decide whether to save extra parameters or not. * Update block.py add moving_mean and moving_var when export model with SymbolBlock * Update python/mxnet/gluon/block.py typo Co-authored-by: Sheng Zha <[email protected]> * Update block.py * Update block.py * Update python/mxnet/gluon/block.py Co-authored-by: Leonard Lausen <[email protected]> Co-authored-by: Sheng Zha <[email protected]> Co-authored-by: Leonard Lausen <[email protected]>
pioy · Sep 1, 2020 · 5122d32 · 5122d32
1 parent 8379740
commit 5122d32
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/python/mxnet/gluon/block.py b/python/mxnet/gluon/block.py
@@ -1348,8 +1348,11 @@ def export(self, path, epoch=0, remove_amp_cast=True):
                 if name in arg_names:
                     arg_dict['arg:{}'.format(name)] = param._reduce()
                 else:
-                    assert name in aux_names
-                    arg_dict['aux:{}'.format(name)] = param._reduce()
+                    if name not in aux_names:
+                        warnings.warn('Parameter "{name}" is not found in the graph. '
+                                      .format(name=name), stacklevel=3)
+                    else:
+                        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)