def distributed_initializer(self):
bcast_ops = []
for v in self.variables():
bcast_ops.append(tf.assign(v, broadcast(v)))
with tf.control_dependencies(bcast_ops):
with tf.control_dependencies([self._save_model_op]):
return barrier()
def test_save_and_request():
global_step = tf.Variable(tf.constant(0, dtype=tf.int64))
target = tf.Variable(tf.constant(0, dtype=tf.int32))
x = tf.Variable(tf.zeros([10], dtype=tf.int32))
inc_op = tf.assign_add(global_step, 1)
update_op = tf.assign(x, x + 1)
save_op = save_variable(x, version=global_step)
y = request_variable(target, global_step, x.name, x.shape, x.dtype)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for i in range(3):
sess.run([inc_op, update_op])
sess.run(save_op)
sess.run(barrier())
v = sess.run(y)
assert v[0] == i + 1
sess.run(barrier())
def test_barrier():
with tf.Session() as sess:
sess.run(barrier())
def init_store(self, variables):
with tf.control_dependencies([self._build_save_op(variables)]):
return barrier()
model = getattr(applications, args.model)(weights=None)
opt = None
learning_rate = 0.01
if args.optimizer == 'sgd':
opt = tf.train.GradientDescentOptimizer(learning_rate)
elif args.optimizer == 'adam':
opt = tf.train.AdamOptimizer(learning_rate)
else:
raise Exception('Unknown optimizer option')
barrier_op = None
if args.kf_optimizer:
from kungfu.tensorflow.ops import barrier
barrier_op = barrier()
if args.kf_optimizer == 'sync-sgd':
from kungfu.tensorflow.optimizers import SynchronousSGDOptimizer
opt = SynchronousSGDOptimizer(opt)
elif args.kf_optimizer == 'sync-sgd-nccl':
from kungfu.tensorflow.optimizers import SynchronousSGDOptimizer
opt = SynchronousSGDOptimizer(opt, nccl=True, nccl_fusion=args.fuse)
elif args.kf_optimizer == 'sync-sgd-hierarchical-nccl':
from kungfu.tensorflow.optimizers import SynchronousSGDOptimizer
opt = SynchronousSGDOptimizer(opt,
nccl=True,
nccl_fusion=args.fuse,
hierarchical_nccl=True)
elif args.kf_optimizer == 'async-sgd':
from kungfu.tensorflow.optimizers import PairAveragingOptimizer
opt = PairAveragingOptimizer(opt, fuse_requests=args.fuse)
