model.py 文件源码

def collect_all_XY(self, root=0):
        if self.mpi_comm is None:
            XY = [self.obslayer.Y.copy()]
            for l in self.layers: XY.append(l.X.copy())
            return XY
        else:
            from mpi4py import MPI
            from GPy.core.parameterization.variational import NormalPosterior
            N,D = self.Y.shape
            N_list = np.array(self.mpi_comm.allgather(N))
            N_all = np.sum(N_list)
            Y_all = np.empty((N_all,D)) if self.mpi_comm.rank==root else None
            self.mpi_comm.Gatherv([self.Y, MPI.DOUBLE], [Y_all, (N_list*D, None), MPI.DOUBLE], root=root)
            if self.mpi_comm.rank==root:
                XY = [Y_all]
            for l in self.layers:
                Q = l.X.shape[1]
                X_mean_all =  np.empty((N_all,Q)) if self.mpi_comm.rank==root else None
                self.mpi_comm.Gatherv([l.X.mean.values, MPI.DOUBLE], [X_mean_all, (N_list*Q, None), MPI.DOUBLE], root=root)
                X_var_all =  np.empty((N_all,Q)) if self.mpi_comm.rank==root else None
                self.mpi_comm.Gatherv([l.X.variance.values, MPI.DOUBLE], [X_var_all, (N_list*Q, None), MPI.DOUBLE], root=root)
                if self.mpi_comm.rank==root:
                    XY.append(NormalPosterior(X_mean_all, X_var_all))
            if self.mpi_comm.rank==root: return XY
            else: return None