def get_updated_variables(loss, current_variables, current_states):
with tf.name_scope("gradients"):
grads = util.get_grads(loss, current_variables)
with tf.name_scope("deltas"):
outputs, next_states, deltas = run_optimizer_for_different_states(
grads, current_states)
next_variables = [x + delta for x, delta in zip(current_variables, deltas)]
return next_variables, list(next_states)
def test_run_optimizer(self):
# we just build the graph here
x = tf.placeholder(dtype=tf.float32, shape=[None, 784])
y = tf.placeholder(dtype=tf.int32, shape=[None, 10])
coordinates, _ = util.get_variables(get_model_loss, {
'x': x,
'y': y,
'custom_variables': None
})
loss = get_model_loss(x, y, custom_variables=coordinates)
grads = util.get_grads(loss, coordinates)
output, state, delta = run_optimizer(grads[0])
self.assertTrue(True)
tf.reset_default_graph()
def test_run_optimizer_different_states(self):
# test if the graph can be built without any
# errors
x = tf.placeholder(dtype=tf.float32, shape=[None, 784])
y = tf.placeholder(dtype=tf.int32, shape=[None, 10])
coordinates, _ = util.get_variables(get_model_loss, {
'x': x,
'y': y,
'custom_variables': None
})
loss = get_model_loss(x, y, custom_variables=coordinates)
grads = util.get_grads(loss, coordinates)
outputs, states, delta = run_optimizer_for_different_states(grads)
self.assertTrue(True)
tf.reset_default_graph()
def test_run_optimizer_build(self):
# we only build the graph, but never run it
x = tf.placeholder(dtype=tf.float32, shape=[None, 784])
y = tf.placeholder(dtype=tf.int32, shape=[None, 10])
coordinates, _ = util.get_variables(get_model_loss, {
'x': x,
'y': y,
'custom_variables': None
})
loss = get_model_loss(x, y, custom_variables=coordinates)
grads = util.get_grads(loss, coordinates)
optimizer = get_optimizer()
hidden = util.get_lstm_state_tuples(optimizer, grads[0])
output, state = optimizer(util.reshape_inputs(grads[0]), hidden)
self.assertTrue(True)
tf.reset_default_graph()
def update_optimizee(self, loss, variables, states):
"""Updates variables by taking the gradient of the
loss w.r.t the variables and run it through the RNN,
then add the output to the variables.
:param loss: optimizee loss
:param variables: an array of tensors containing the past
optimizee variables
:param states: LSTM states
:return: tuple containing new variables and list of new
LSTM states
"""
with tf.name_scope("gradients"):
grads = util.get_grads(loss, variables)
with tf.name_scope("deltas"):
_, new_states, deltas = self.run_multiple_states(grads, states)
with tf.name_scope("add_deltas"):
new_variables = [x + delta for x, delta in zip(variables, deltas)]
return new_variables, list(new_states)
def unroll(input_arguments):
"""Creates the unroll operations and returns the final loss,
variables and aditional update and reset operations used
to start a new training epoch.
:param input_arguments: dictionary containing arguments for
optimizee loss
:returns: a tuple containing the final loss (reduce sum),
reset and update operations, the final coordinates
(or variables) and the final loss operation
"""
optimizer = Optimizer()
# get list of variables and constants
fake_optee = OptimizeeFC()
variables, constants = util.get_vars(fake_optee.loss,
{'inputs': tf.placeholder(tf.float32,
shape=[None, 784]),
'labels': tf.placeholder(tf.int32,
shape=[None, 10])},
scope=MODEL_CONSTANTS.FINAL_SCOPE)
# get and reshape RNN states for each variable
with tf.name_scope("states"):
reshaped_variables = [util.reshape_inputs(x) for x in variables]
initial_states = [optimizer.get_input_states(
x) for x in reshaped_variables]
# initialize array for keeping track of the loss values in loomp
loss_array = tf.TensorArray(tf.float32, size=FLAGS.truncated_backprop+1,
clear_after_read=False)
# run loop and collect the loss values in the array
_, _, loss_array, new_variables, new_states = tf.while_loop(
cond=lambda t, *_: t < FLAGS.truncated_backprop,
body=loop_body,
loop_vars=(0, input_arguments, loss_array, variables, initial_states),
swap_memory=True,
parallel_iterations=1,
name="unroll")
# run last time without any gradient update
with tf.name_scope("final_loss"):
fl_optimizee = OptimizeeFC()
random_x, random_y = random_batch()
args = {'inputs': random_x, 'labels': random_y}
final_loss = fl_optimizee.loss(
custom_variables=new_variables, **args)
loss_array = loss_array.write(FLAGS.truncated_backprop, final_loss)
# print some logs: the gradient histogram so we see how they
# evolved and the difference in vectors if they are adversarial
if FLAGS.instrumentation:
grads = util.get_grads(final_loss, new_variables)
util.save_grad_mean(grads)
if FLAGS.adversarial:
adv_ex = fl_optimizee.generate_adv(random_x)
_, diff = util.is_same(adv_ex, random_x)
tf.summary.scalar('vector-difference', diff)
# collect the final loss and print it to file
sum_loss = tf.reduce_sum(loss_array.stack(), name="optimizee_sum_loss")
if FLAGS.instrumentation:
tf.summary.scalar("loss/final_loss", final_loss)
tf.summary.scalar("loss/average_loss", sum_loss)
# create reset op to be called at the begining of each new epoch
with tf.name_scope("reset_loop_vars"):
reset_variables = (nest.flatten(util.nested_variable(
initial_states)) + variables + constants)
reset_op = [tf.variables_initializer(reset_variables), loss_array.close()]
# create update op to update the final vars
with tf.name_scope("update"):
update_op = (nest.flatten(
util.nested_assign(variables, new_variables)) +
nest.flatten(util.nested_assign(
util.nested_variable(initial_states),
new_states)))
optimize_op = optimizer.optimize(sum_loss)
return optimize_op, reset_op, update_op, final_loss, new_variables