Skip to content

Instantly share code, notes, and snippets.

Last active February 25, 2019 13:07
What would you like to do?
# source:
with tf.device('/gpu:0'), \
data, target, loss = create_simple_model(nbatch, nin, nout, dtype)
variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# Traning variables
lr = 0.0002
beta = 0.5
variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
model_opt = tf.train.AdamOptimizer(learning_rate, beta) # Adam optimizer
# Mixed precision
scale_size = 128 # There is no one scale size
loss_scale_manager = FixedLossScaleManager(scale)
loss_scale_optimizer = LossScaleOptimizer(model_opt, loss_scale_manager)
# Calculate the gradients with scaled loss and return the unscaled gradients
grads_variables = loss_scale_optimizer.compute_gradients(loss, variables)
Doing some gradient manipulation (if needed)
only example!
grads_variables = [(g,v) for (g,v) in grads_variables if g is not None]
training_opt = loss_scale_optimizer.apply_gradients(grads_variables)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment