gaspersic/baselines
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

use allreduce instead of Allreduce (send pickled data instead of floats) - probably affects performance somewhat, but avoid element number mismatch. Fixes 998 (#1000)

Browse Source
This commit is contained in:
pzhokhov
2019-09-27 14:45:31 -07:00
committed by GitHub
parent 229a772b81
commit 53797293e5
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
baselines/common/mpi_moments.py
View File

@@ -12,8 +12,9 @@ def mpi_mean(x, axis=0, comm=None, keepdims=False):
localsum = np.zeros(n+1, x.dtype) localsum = np.zeros(n+1, x.dtype)
localsum[:n] = xsum.ravel() localsum[:n] = xsum.ravel()
localsum[n] = x.shape[axis] localsum[n] = x.shape[axis]
globalsum = np.zeros_like(localsum) # globalsum = np.zeros_like(localsum)
comm.Allreduce(localsum, globalsum, op=MPI.SUM) # comm.Allreduce(localsum, globalsum, op=MPI.SUM)
globalsum = comm.allreduce(localsum, op=MPI.SUM)
return globalsum[:n].reshape(xsum.shape) / globalsum[n], globalsum[n] return globalsum[:n].reshape(xsum.shape) / globalsum[n], globalsum[n]
def mpi_moments(x, axis=0, comm=None, keepdims=False): def mpi_moments(x, axis=0, comm=None, keepdims=False):