26 lines
556 B
Python
Executable File
26 lines
556 B
Python
Executable File
#!/usr/bin/python
|
|
import rabit
|
|
import numpy as np
|
|
|
|
rabit.init(lib='mock')
|
|
rank = rabit.get_rank()
|
|
n = 10
|
|
nround = 3
|
|
data = np.ones(n) * rank
|
|
|
|
version, model, local = rabit.load_checkpoint(True)
|
|
if version == 0:
|
|
model = np.zeros(n)
|
|
local = np.ones(n)
|
|
else:
|
|
print '[%d] restart from version %d' % (rank, version)
|
|
|
|
for i in xrange(version, nround):
|
|
res = rabit.allreduce(data + model+local, rabit.SUM)
|
|
print '[%d] iter=%d: %s' % (rank, i, str(res))
|
|
model = res
|
|
local[:] = i
|
|
rabit.checkpoint(model, local)
|
|
|
|
rabit.finalize()
|