Optimizers

MP_optimizer.py

# source: https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html

with tf.device('/gpu:0'), \
    tf.variable_scope('fp32_storage',custom_getter=float32_variable_storage_getter):
        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)