From 9799bce36f99ba9c020c2f009328292bca5d23c2 Mon Sep 17 00:00:00 2001 From: AshwinParanjape Date: Wed, 16 May 2018 14:18:03 -0700 Subject: [PATCH 1/2] Faster testing code * Creating the softmax_W tensor on CPU makes it transfer the tensor from CPU to GPU for every batch during test time (where full softmax is used), slowing down the code by 5x or so. * Transposing softmax_W is a large operation compared to transposing input and output vectors, increasing the speed of full softmax further --- bilm/training.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/bilm/training.py b/bilm/training.py index 2d60951..890447a 100644 --- a/bilm/training.py +++ b/bilm/training.py @@ -461,7 +461,7 @@ def _get_next_token_placeholders(suffix): # softmax_W is just the embedding layer self.softmax_W = self.embedding_weights - with tf.variable_scope('softmax'), tf.device('/cpu:0'): + with tf.variable_scope('softmax'): # Glorit init (std=(1.0 / sqrt(fan_in)) softmax_init = tf.random_normal_initializer(0.0, 1.0 / np.sqrt(softmax_dim)) @@ -502,10 +502,10 @@ def _get_next_token_placeholders(suffix): else: # get the full softmax loss - output_scores = tf.matmul( - lstm_output_flat, - tf.transpose(self.softmax_W) - ) + self.softmax_b + output_scores = tf.transpose(tf.matmul( + self.softmax_W, + tf.transpose(lstm_output_flat) + )) + self.softmax_b # NOTE: tf.nn.sparse_softmax_cross_entropy_with_logits # expects unnormalized output since it performs the # softmax internally From c307638226ba9256ab575533f63607c6bc6ca40c Mon Sep 17 00:00:00 2001 From: AshwinParanjape Date: Wed, 16 May 2018 17:03:20 -0700 Subject: [PATCH 2/2] CPU placement of softmax_W during training --- bilm/training.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bilm/training.py b/bilm/training.py index 890447a..603292f 100644 --- a/bilm/training.py +++ b/bilm/training.py @@ -461,7 +461,7 @@ def _get_next_token_placeholders(suffix): # softmax_W is just the embedding layer self.softmax_W = self.embedding_weights - with tf.variable_scope('softmax'): + with tf.variable_scope('softmax'), tf.device(lambda op: '' if not self.is_training else '/cpu:0'): # Glorit init (std=(1.0 / sqrt(fan_in)) softmax_init = tf.random_normal_initializer(0.0, 1.0 / np.sqrt(softmax_dim))