apeforest commented on a change in pull request #12666: [WIP] Horovod support
for MXNet
URL: https://github.com/apache/incubator-mxnet/pull/12666#discussion_r226113346
##########
File path: python/mxnet/module/executor_group.py
##########
@@ -421,17 +421,31 @@ def get_params(self, arg_params, aux_params):
Target parameter arrays.
aux_params : list of NDArray
Target aux arrays.
+ copy_to_cpu : boolean
+ Whether or not to copy parameters to CPU. (default to 'true')
Notes
-----
- This function will inplace update the NDArrays in arg_params and
aux_params.
"""
- for name, block in zip(self.param_names, self.param_arrays):
- weight = sum(w.copyto(ctx.cpu()) for w in block) / len(block)
- weight.astype(arg_params[name].dtype).copyto(arg_params[name])
- for name, block in zip(self.aux_names, self.aux_arrays):
- weight = sum(w.copyto(ctx.cpu()) for w in block) / len(block)
- weight.astype(aux_params[name].dtype).copyto(aux_params[name])
+ if copy_to_cpu:
+ for name, block in zip(self.param_names, self.param_arrays):
+ weight = sum(w.copyto(ctx.cpu()) for w in block) / len(block)
+ weight.astype(arg_params[name].dtype).copyto(arg_params[name])
+ for name, block in zip(self.aux_names, self.aux_arrays):
+ weight = sum(w.copyto(ctx.cpu()) for w in block) / len(block)
+ weight.astype(aux_params[name].dtype).copyto(aux_params[name])
+ else:
+ for name, block in zip(self.param_names, self.param_arrays):
+ context = block[0].context
Review comment:
It seems this `else` block is identical to the `if` block except for the
context. Can you add the if check to the context and change context between
`block[0].context` and `ctx.cpu()`
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
