daily pastebin goal
57%
SHARE
TWEET

Untitled

a guest Mar 22nd, 2019 62 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import tensorflow as tf
  2. import numpy as np
  3. def float32_variable_storage_getter(getter,name,shape=None,dtype=None,
  4.                                     initializer=None, regularizer=None,
  5.                                     trainable = True,
  6.                                     *args,**kwargs):
  7.     """
  8.     Custom variable getter that forces trainable variables to be stored in
  9.     float32 precision and then casts them to the training precision.
  10.     """
  11.     storage_dtype = tf.float32 if trainable else dtype
  12.     variable = getter(name,shape,dtype=storage_dtype,
  13.                       initializer =initializer, regularizer = regularizer,
  14.                       *args, **kwargs)
  15.                      
  16.     if trainable and dtype != tf.float32:
  17.          variable = tf.cast(variable,dtype)
  18.     return variable
  19. def gradients_with_loss_scaling(loss,variables,loss_scale):
  20.     """Gradient calculation with loss scaling to improve numerical stability when training with float16.
  21.     """
  22.     return [grad/loss_scale for grad in tf.gradients(loss*loss_scale,variables)]
  23.  
  24. def create_simple_model(nbatch, nin, nout, dtype):
  25.     """A simple softmax model."""
  26.     data    = tf.placeholder(dtype, shape=(nbatch, nin))
  27.     weights = tf.get_variable('weights', (nin, nout), dtype)
  28.     biases  = tf.get_variable('biases', nout, dtype,initializer=tf.zeros_initializer())
  29.    
  30.     logits  = tf.matmul(data, weights) + biases
  31.    
  32.     target  = tf.placeholder(tf.float32, shape=(nbatch, nout))
  33.     # Note: The softmax should be computed in float32 precision
  34.     loss    = tf.losses.softmax_cross_entropy(target, tf.cast(logits, tf.float32))
  35.     return data, target, loss
  36.    
  37. if __name__ == '__main__':
  38.     nbatch= 64
  39.     nin=104
  40.     nout = 16
  41.     learning_rate = 0.1
  42.     momentum = 0.9
  43.     loss_scale = 128
  44.     dtype = tf.float16
  45.     tf.set_random_seed(1234)
  46.     np.random.seed(4321)
  47.    
  48.     #Create training graph
  49.     with tf.device('/gpu:0'), \
  50.          tf.variable_scope(
  51.              # Note: This forces trainable variables to be stored as float32
  52.              'fp32_storage', custom_getter=float32_variable_storage_getter):
  53.         data, target, loss = create_simple_model(nbatch, nin, nout, dtype)
  54.         variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
  55.         # Note: Loss scaling can improve numerical stability for fp16 training
  56.         grads = gradients_with_loss_scaling(loss, variables, loss_scale)
  57.         optimizer = tf.train.MomentumOptimizer(learning_rate, momentum)
  58.         training_step_op = optimizer.apply_gradients(zip(grads, variables))
  59.         init_op = tf.global_variables_initializer()
  60.  
  61.     # Run training
  62.     sess = tf.Session()
  63.     sess.run(init_op)
  64.     np_data   = np.random.normal(size=(nbatch, nin)).astype(np.float16)
  65.     np_target = np.zeros((nbatch, nout), dtype=np.float32)
  66.     np_target[:,0] = 1
  67.     print 'Step Loss'
  68.     for step in xrange(30):
  69.         np_loss, _ = sess.run([loss, training_step_op],
  70.                               feed_dict={data: np_data, target: np_target})
  71.         print '%4i %6f' % (step + 1, np_loss)
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top