def run(self, inputs, states=None):
"""Runs the optimizer once for given
inputs and states and returns the results.
:param inputs: a tensor containing the inputs / gradients
:param states: an array containing LSTMTuples
"""
self.check_cells()
if states is None:
states = self.get_input_states(inputs)
with tf.name_scope("preprocess_rnn"):
reshaped_inputs = util.reshape_inputs(inputs)
preprocessed = util.preprocess(tf.expand_dims(reshaped_inputs, -1))
# adds preprocessing data and dimension
reshaped_inputs = tf.reshape(
preprocessed, [preprocessed.get_shape().as_list()[0], -1])
with tf.name_scope("run_rnn"):
output, new_state = self.cells(reshaped_inputs, states)
delta = self.get_linear(output)
with tf.name_scope("postprocess_rnn"):
delta = delta * FLAGS.optimizer_scale # add rescaling
reshaped = util.recover_input_shape(delta, inputs)
return output, new_state, reshaped
def run_optimizer(inputs, states=None):
optimizer = get_optimizer()
reshaped_inputs = util.reshape_inputs(inputs)
if states is None:
states = util.get_lstm_state_tuples(optimizer, reshaped_inputs)
output, new_state = optimizer(reshaped_inputs, states)
delta = get_linear(output, 1)
return output, new_state, util.recover_input_shape(delta, inputs)
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 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
def get_input_states(optimizer, inputs):
reshaped_inputs = [util.reshape_inputs(x) for x in inputs]
return [util.get_lstm_state_tuples(optimizer, x) for x in reshaped_inputs]