def test_dist():
np.random.seed(0)
p1,p2,p3=(np.random.randn(3,1), np.random.randn(4,1), np.random.randn(5,1))
q1,q2,q3=(np.random.randn(6,1), np.random.randn(7,1), np.random.randn(8,1))
# p1,p2,p3=(np.random.randn(3), np.random.randn(4), np.random.randn(5))
# q1,q2,q3=(np.random.randn(6), np.random.randn(7), np.random.randn(8))
comm = MPI.COMM_WORLD
assert comm.Get_size()==2
if comm.Get_rank()==0:
x1,x2,x3 = p1,p2,p3
elif comm.Get_rank()==1:
x1,x2,x3 = q1,q2,q3
else:
assert False
rms = RunningMeanStd(epsilon=0.0, shape=(1,))
U.initialize()
rms.update(x1)
rms.update(x2)
rms.update(x3)
bigvec = np.concatenate([p1,p2,p3,q1,q2,q3])
def checkallclose(x,y):
print(x,y)
return np.allclose(x,y)
assert checkallclose(
bigvec.mean(axis=0),
U.eval(rms.mean)
)
assert checkallclose(
bigvec.std(axis=0),
U.eval(rms.std)
)
python类COMM_WORLD的实例源码
def __init__(self, var_list, *, beta1=0.9, beta2=0.999, epsilon=1e-08, scale_grad_by_procs=True, comm=None):
self.var_list = var_list
self.beta1 = beta1
self.beta2 = beta2
self.epsilon = epsilon
self.scale_grad_by_procs = scale_grad_by_procs
size = sum(U.numel(v) for v in var_list)
self.m = np.zeros(size, 'float32')
self.v = np.zeros(size, 'float32')
self.t = 0
self.setfromflat = U.SetFromFlat(var_list)
self.getflat = U.GetFlat(var_list)
self.comm = MPI.COMM_WORLD if comm is None else comm
def run(self, clf):
"""Run correlation-based voxel selection in master-worker model.
Sort the voxels based on the cross-validation accuracy
of their correlation vectors
Parameters
----------
clf: classification function
the classifier to be used in cross validation
Returns
-------
results: list of tuple (voxel_id, accuracy)
the accuracy numbers of all voxels, in accuracy descending order
the length of array equals the number of voxels
"""
rank = MPI.COMM_WORLD.Get_rank()
if rank == self.master_rank:
results = self._master()
# Sort the voxels
results.sort(key=lambda tup: tup[1], reverse=True)
else:
self._worker(clf)
results = []
return results
def _worker(self, clf):
"""Worker node's operation.
Receiving tasks from the master to process and sending the result back
Parameters
----------
clf: classification function
the classifier to be used in cross validation
Returns
-------
None
"""
logger.debug(
'worker %d is running, waiting for tasks from master at rank %d' %
(MPI.COMM_WORLD.Get_rank(), self.master_rank)
)
comm = MPI.COMM_WORLD
status = MPI.Status()
while 1:
task = comm.recv(source=self.master_rank,
tag=MPI.ANY_TAG,
status=status)
if status.Get_tag():
break
comm.send(self._voxel_scoring(task, clf),
dest=self.master_rank)
def __init__(self, sl_rad=1, max_blk_edge=10, shape=Cube,
min_active_voxels_proportion=0):
"""Constructor
Parameters
----------
sl_rad: radius, in voxels, of the sphere inscribed in the
searchlight cube, not counting the center voxel
max_blk_edge: max edge length, in voxels, of the 3D block
shape: brainiak.searchlight.searchlight.Shape indicating the
shape in voxels of the searchlight region
min_active_voxels_proportion: float
If a searchlight region does not have more than this minimum
proportion of active voxels in the mask, it is not processed by the
searchlight function. The mask used for the test is the
intersection of the global (brain) mask and the `Shape` mask. The
seed (central) voxel of the searchlight region is taken into
consideration.
"""
self.sl_rad = sl_rad
self.max_blk_edge = max_blk_edge
self.min_active_voxels_proportion = min_active_voxels_proportion
self.comm = MPI.COMM_WORLD
self.shape = shape(sl_rad).mask_
self.bcast_var = None
def test_searchlight_with_cube():
sl = Searchlight(sl_rad=3)
comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size
dim0, dim1, dim2 = (50, 50, 50)
ntr = 30
nsubj = 3
mask = np.zeros((dim0, dim1, dim2), dtype=np.bool)
data = [np.empty((dim0, dim1, dim2, ntr), dtype=np.object)
if i % size == rank
else None
for i in range(0, nsubj)]
# Put a spot in the mask
mask[10:17, 10:17, 10:17] = True
sl.distribute(data, mask)
global_outputs = sl.run_searchlight(cube_sfn)
if rank == 0:
assert global_outputs[13, 13, 13] == 1.0
global_outputs[13, 13, 13] = None
for i in range(global_outputs.shape[0]):
for j in range(global_outputs.shape[1]):
for k in range(global_outputs.shape[2]):
assert global_outputs[i, j, k] is None
def test_searchlight_with_diamond():
sl = Searchlight(sl_rad=3, shape=Diamond)
comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size
dim0, dim1, dim2 = (50, 50, 50)
ntr = 30
nsubj = 3
mask = np.zeros((dim0, dim1, dim2), dtype=np.bool)
data = [np.empty((dim0, dim1, dim2, ntr), dtype=np.object)
if i % size == rank
else None
for i in range(0, nsubj)]
# Put a spot in the mask
mask[10:17, 10:17, 10:17] = Diamond(3).mask_
sl.distribute(data, mask)
global_outputs = sl.run_searchlight(diamond_sfn)
if rank == 0:
assert global_outputs[13, 13, 13] == 1.0
global_outputs[13, 13, 13] = None
for i in range(global_outputs.shape[0]):
for j in range(global_outputs.shape[1]):
for k in range(global_outputs.shape[2]):
assert global_outputs[i, j, k] is None
def __init__(self, fortran=None, *args, **kwargs):
"""
To use the pyOM2 legacy interface point the fortran argument to the Veros fortran library:
> simulation = GlobalOneDegree(fortran = "pyOM_code.so")
"""
if fortran:
self.legacy_mode = True
try:
self.fortran = LowercaseAttributeWrapper(imp.load_dynamic("pyOM_code", fortran))
self.use_mpi = False
except ImportError:
self.fortran = LowercaseAttributeWrapper(imp.load_dynamic("pyOM_code_MPI", fortran))
self.use_mpi = True
from mpi4py import MPI
self.mpi_comm = MPI.COMM_WORLD
self.main_module = LowercaseAttributeWrapper(self.fortran.main_module)
self.isoneutral_module = LowercaseAttributeWrapper(self.fortran.isoneutral_module)
self.idemix_module = LowercaseAttributeWrapper(self.fortran.idemix_module)
self.tke_module = LowercaseAttributeWrapper(self.fortran.tke_module)
self.eke_module = LowercaseAttributeWrapper(self.fortran.eke_module)
else:
self.legacy_mode = False
self.use_mpi = False
self.fortran = self
self.main_module = self
self.isoneutral_module = self
self.idemix_module = self
self.tke_module = self
self.eke_module = self
self.modules = (self.main_module, self.isoneutral_module, self.idemix_module,
self.tke_module, self.eke_module)
if self.use_mpi and self.mpi_comm.Get_rank() != 0:
kwargs["loglevel"] = "critical"
super(VerosLegacy, self).__init__(*args, **kwargs)
def main(tiff, outfile, overlayname):
"""
Turn a geotiff into a KMZ that can be dragged onto an instance of Terria
Map. This also constructs a JPEG of the Geotiff, as it is required for the
KMZ.
"""
# MPI globals
comm = MPI.COMM_WORLD
chunk_index = comm.Get_rank()
# This runs on the root node only
if chunk_index != 0:
return
# Get tiff info
I = geoio.Image(tiff)
# Save tiff as jpeg
if outfile is not None:
outfile = os.path.splitext(outfile)[0]
else:
outfile = os.path.splitext(tiff)[0]
jpg = outfile + ".jpg"
# Convert tiff to jpg
Im = Image.open(tiff)
Im.save(jpg)
# Construct KMZ
kml = simplekml.Kml()
if overlayname is None:
overlayname = os.path.basename(outfile)
ground = kml.newgroundoverlay(name=overlayname)
ground.icon.href = jpg
ground.latlonbox.west = I.xmin
ground.latlonbox.east = I.xmax
ground.latlonbox.north = I.ymax
ground.latlonbox.south = I.ymin
kml.savekmz("{}.kmz".format(outfile))
def convert_files(files, output_dir, opt, mask_file):
comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size
if mask_file:
# temporary cropped mask file
cropped_mask_file = tempfile.mktemp(suffix='.tif', dir=TMPDIR)
# crop/reproject/resample the mask
crop_reproject_resample(opt.mask_file,
cropped_mask_file,
sampling='near',
extents=opt.extents,
reproject=opt.reproject)
else:
cropped_mask_file = opt.mask_file
for i in range(rank, len(files), size):
in_file = files[i]
print('============file no: {} of {}==========='.format(i, len(files)))
log.info("operating on {file} in process {rank}".format(file=in_file,
rank=rank))
out_file = join(output_dir, basename(in_file))
log.info('Output file: {}'.format(out_file))
do_work(input_file=in_file,
mask_file=cropped_mask_file,
output_file=out_file,
options=options)
if mask_file:
os.remove(cropped_mask_file)
log.info('removed intermediate cropped '
'mask file {}'.format(cropped_mask_file))
comm.Barrier()
# TODO: use click here
def write_server_info(filename, port):
pid = os.getpid()
rank = MPI.COMM_WORLD.Get_rank()
server_info = '{}:{}:{}:{}:{}'.format(LINE_TOKEN, rank, pid, port, LINE_TOKEN).strip()
logger.debug("write_server_info: line %s, filename %s", server_info, filename)
time.sleep(0.1 * rank)
with open(filename, "a") as f:
fcntl.lockf(f, fcntl.LOCK_EX)
f.write(server_info + '\n')
f.flush()
os.fsync(f.fileno())
fcntl.lockf(f, fcntl.LOCK_UN)
return server_info
def serve():
parser = argparse.ArgumentParser()
parser.add_argument("-tf", "--tmpfile", nargs=1)
parser.add_argument("-p", "--ports", nargs='+')
args = parser.parse_args()
comm = MPI.COMM_WORLD
options = [('grpc.max_send_message_length', -1), ('grpc.max_receive_message_length', -1)]
server = grpc.server(futures.ThreadPoolExecutor(max_workers=1), options=options)
hetr_pb2_grpc.add_HetrServicer_to_server(HetrServer(comm, server), server)
logger.debug("server: rank %d, tmpfile %s, ports %s",
comm.Get_rank(), args.tmpfile[0], args.ports if args.ports is not None else "")
if args.ports is not None and len(args.ports) > comm.Get_rank():
p = args.ports[comm.Get_rank()]
if is_port_open(p):
port = server.add_insecure_port('[::]:' + p)
else:
raise RuntimeError("port %s is already in use!", p)
else:
port = server.add_insecure_port('[::]:0')
server.start()
write_server_info(args.tmpfile[0], port)
try:
while True:
time.sleep(_ONE_DAY_IN_SECONDS)
except KeyboardInterrupt:
server.stop(0)
def __init__(self, target=None, comm=None, debug=False, loadbalance=False):
self.comm = MPI.COMM_WORLD if comm is None else comm
self.rank = self.comm.Get_rank()
self.size = self.comm.Get_size() - 1
self.debug = debug
self.function = _error_function if target is None else target
self.loadbalance = loadbalance
if self.size == 0:
raise ValueError("Tried to create an MPI pool, but there "
"was only one MPI process available. "
"Need at least two.")
def get_internode_comm(self):
self.comm=MPI.COMM_WORLD
self.rank=self.comm.rank
self.size=self.comm.size
def get_internode_comm():
from mpi4py import MPI
comm=MPI.COMM_WORLD
return comm
# intra-node comm
def get_internode_comm():
from mpi4py import MPI
comm=MPI.COMM_WORLD
return comm
# intra-node comm
def bcast(self, data):
if MPI_INSTALLED:
mpi_comm = MPI.COMM_WORLD
bdata = mpi_comm.bcast(data, root=0)
else:
bdata = data
return bdata
# Wrapper for common MPI interfaces.
def barrier(self):
if MPI_INSTALLED:
mpi_comm = MPI.COMM_WORLD
mpi_comm.barrier()
def rank(self):
if MPI_INSTALLED:
mpi_comm = MPI.COMM_WORLD
return mpi_comm.Get_rank()
else:
return 0
def size(self):
if MPI_INSTALLED:
mpi_comm = MPI.COMM_WORLD
return mpi_comm.Get_size()
else:
return 1
