def _build_request_and_save_ops(self, target, variables):
var_fused = fuse(variables)
save_model_op = save_variable(var_fused)
other_peer_var_fused = request_variable_with_template(
target, var_fused)
other_peer_vars = defuse(other_peer_var_fused,
[v.shape for v in variables])
self._save_model_op = save_model_op # save for _get_initializer_op
return other_peer_vars, save_model_op
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 _build_save_op(self, variables):
if self._fuse_requests:
var_fused = fuse(variables)
return save_variable(var_fused, name=self._fused_model_name)
else:
return tf.group([save_variable(v) for v in variables])
