leezu commented on a change in pull request #17970:
URL: https://github.com/apache/incubator-mxnet/pull/17970#discussion_r480249528
##########
File path: python/mxnet/gluon/block.py
##########
@@ -1343,13 +1343,18 @@ def export(self, path, epoch=0, remove_amp_cast=True):
arg_names = set(sym.list_arguments())
aux_names = set(sym.list_auxiliary_states())
arg_dict = {}
- for is_arg, name, param in self._cached_op_args:
- if not is_arg:
- if name in arg_names:
- arg_dict['arg:{}'.format(name)] = param._reduce()
- else:
+ for name, param in self.collect_params().items():
+ if name in arg_names:
+ arg_dict['arg:%s'%name] = param._reduce()
+ else:
+ 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:{}'.format(name)] = param._reduce()
+ arg_dict['aux:%s'%name] = param._reduce()
+ else:
+ warnings.warn('Parameter "{name}" is ignored in saving '
+ 'as it is not found in the graph. '
+ .format(name=name), stacklevel=3)
Review comment:
This will ignore all aux parameters that are part of the graph, but are
not using the batch norm naming scheme. Why not just remove the `assert name in
aux_names` from the original implementation and save all auxiliary parameters
that are used in the graph? I don't see a reason for restricting the set of
saved parameters in the frontend API.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
