p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!9550 Optimize performance of PyNative grad reduce 

From: @jojobugfree
Reviewed-by: @kisnwang,@jjfeing
Signed-off-by: @jjfeing
This commit is contained in:
mindspore-ci-bot 2020-12-05 16:16:56 +08:00 committed by Gitee
parent edb145f8f6 23898c735f
commit e4f1365495
1 changed files with 1 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mindspore/nn/wrap/grad_reducer.py
View File

@ -65,9 +65,7 @@ def _tensors_allreduce(degree, mean, allgather, allreduce, allreduce_filter, gra
grad = allreduce(grad)
if mean:
degree = F.scalar_cast(degree, F.dtype(grad))
cast_op = P.Cast()
mul_op = P.Mul()
grad = mul_op(grad, cast_op(F.scalar_to_array(1.0 / degree), F.dtype(grad)))
grad = F.tensor_mul(grad, F.cast(F.scalar_to_array(1.0 / degree), F.dtype(grad)))
return grad
return grad