shape[layer] = [size/(layer+1)]
total_size += shape[layer][0]
src_bufs = [None]*nlayer
src_backups = [None]*nlayer
bufs_expect = [None]*nlayer
for layer in range(nlayer):
src_bufs[layer] = numpy.zeros(shape[layer], numpy.float32)
src_backups[layer] = numpy.zeros(shape[layer], numpy.float32)
bufs_expect[layer] = numpy.zeros(shape[layer], numpy.float32)
distribute.init(6)
world_size = distribute.get_world_size()
rank = distribute.get_rank()
for layer in range(nlayer):
src_bufs[layer] = (
numpy.full(shape[layer], rank+layer*10, numpy.float32)
+ numpy.linspace(0.0,
(shape[layer][0]+0.0)/(shape[layer][0]+1.0),
num=shape[layer][0], endpoint=False,
dtype=numpy.float32))
for layer in range(nlayer):
src_bufs[layer] = ideep4py.mdarray(src_bufs[layer])
src_backups[layer] = ideep4py.mdarray(src_backups[layer])
ideep4py.basic_copyto(src_backups[layer], src_bufs[layer])
iter_num = 50
def cb(id):
print("[%d]callback from payload %d" % (distribute.get_rank(), id))