| # Copyright 2016 The TensorFlow Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # ============================================================================== |
| """Tests for Cudnn RNN models.""" |
| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
| |
| import argparse |
| import collections |
| import functools |
| import itertools |
| import os |
| import sys |
| import unittest |
| |
| import numpy as np |
| |
| from tensorflow.contrib.cudnn_rnn.python.layers import cudnn_rnn |
| from tensorflow.contrib.cudnn_rnn.python.ops import cudnn_rnn_ops |
| from tensorflow.contrib.rnn.python.ops import rnn as contrib_rnn_lib |
| from tensorflow.python.eager import backprop |
| from tensorflow.python.eager import context |
| from tensorflow.python.framework import dtypes |
| from tensorflow.python.framework import ops |
| from tensorflow.python.framework import random_seed |
| from tensorflow.python.framework import test_util |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.ops import control_flow_ops |
| from tensorflow.python.ops import gen_nn_ops |
| from tensorflow.python.ops import gradients_impl as gradients |
| from tensorflow.python.ops import init_ops |
| from tensorflow.python.ops import math_ops |
| from tensorflow.python.ops import random_ops |
| from tensorflow.python.ops import rnn as rnn_lib |
| from tensorflow.python.ops import rnn_cell_impl |
| from tensorflow.python.ops import state_ops |
| from tensorflow.python.ops import variable_scope as vs |
| from tensorflow.python.ops import variables |
| from tensorflow.python.ops.losses import losses |
| from tensorflow.python.platform import googletest |
| from tensorflow.python.platform import test |
| from tensorflow.python.platform import tf_logging as logging |
| from tensorflow.python.training import adagrad |
| from tensorflow.python.training import adam |
| from tensorflow.python.training import gradient_descent |
| from tensorflow.python.training import momentum |
| from tensorflow.python.training import rmsprop |
| from tensorflow.python.training import saver as saver_lib |
| from tensorflow.python.training.checkpointable import util as checkpointable_utils |
| |
| |
| CUDNN_LSTM = cudnn_rnn_ops.CUDNN_LSTM |
| CUDNN_GRU = cudnn_rnn_ops.CUDNN_GRU |
| CUDNN_RNN_RELU = cudnn_rnn_ops.CUDNN_RNN_RELU |
| CUDNN_RNN_TANH = cudnn_rnn_ops.CUDNN_RNN_TANH |
| CUDNN_RNN_UNIDIRECTION = cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION |
| CUDNN_RNN_BIDIRECTION = cudnn_rnn_ops.CUDNN_RNN_BIDIRECTION |
| |
| CUDNN_LSTM_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_LSTM_PARAMS_PER_LAYER |
| CUDNN_GRU_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_GRU_PARAMS_PER_LAYER |
| CUDNN_RNN_TANH_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_RNN_TANH_PARAMS_PER_LAYER |
| CUDNN_RNN_RELU_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_RNN_RELU_PARAMS_PER_LAYER |
| |
| |
| class CudnnTestModel(object): |
| """Model with convenient APIs for easier building and running test graph. |
| |
| The graph built is used by all tests below to avoid repeatedly building |
| similar test graphs. |
| """ |
| |
| def __init__(self, |
| rnn_mode, |
| num_layers, |
| num_units, |
| input_size, |
| direction=CUDNN_RNN_UNIDIRECTION, |
| dropout=0., |
| dtype=dtypes.float32, |
| training=False, |
| seed=None, |
| kernel_initializer=None, |
| bias_initializer=None): |
| if dtype not in (dtypes.float16, dtypes.float32, dtypes.float64): |
| raise ValueError("Invalid dtype: %s" % dtype) |
| self._dtype = dtype |
| |
| self._inputs = array_ops.placeholder( |
| dtype=dtype, shape=[None, None, input_size], name="inputs") |
| h = array_ops.placeholder( |
| dtype=dtype, shape=[None, None, num_units], name="h") |
| c = array_ops.placeholder( |
| dtype=dtype, shape=[None, None, num_units], name="c") |
| if rnn_mode == CUDNN_LSTM: |
| model_fn = cudnn_rnn.CudnnLSTM |
| self._initial_state = (h, c) |
| elif rnn_mode == CUDNN_GRU: |
| model_fn = cudnn_rnn.CudnnGRU |
| self._initial_state = (h,) |
| elif rnn_mode == CUDNN_RNN_TANH: |
| model_fn = cudnn_rnn.CudnnRNNTanh |
| self._initial_state = (h,) |
| elif rnn_mode == CUDNN_RNN_RELU: |
| model_fn = cudnn_rnn.CudnnRNNRelu |
| self._initial_state = (h,) |
| else: |
| raise ValueError("Invalid rnn_mode: %s" % rnn_mode) |
| self._rnn = model_fn( |
| num_layers, |
| num_units, |
| direction=direction, |
| dropout=dropout, |
| dtype=dtype, |
| seed=seed, |
| kernel_initializer=kernel_initializer, |
| bias_initializer=bias_initializer) |
| self._rnn.build([None, None, input_size]) |
| |
| self._outputs, self._output_state = self._rnn( |
| self._inputs, initial_state=self._initial_state, training=training) |
| |
| def _AddUp(self, outputs, output_state): |
| total = math_ops.reduce_sum(outputs) |
| for s in output_state: |
| total += math_ops.reduce_sum(s) |
| return total |
| |
| @property |
| def inputs(self): |
| return self._inputs |
| |
| @property |
| def initial_state(self): |
| return self._initial_state |
| |
| @property |
| def outputs(self): |
| return self._outputs |
| |
| @property |
| def output_state(self): |
| return self._output_state |
| |
| @property |
| def rnn(self): |
| return self._rnn |
| |
| @property |
| def total_sum(self): |
| return self._AddUp(self.outputs, self.output_state) |
| |
| def SynthesizeInput(self, seq_length, batch_size, seed=1234): |
| """Synthesizes input and initial state values for testing.""" |
| np.random.seed(seed) |
| num_layers = self._rnn.num_layers |
| dir_count = self._rnn.num_dirs |
| num_units = self._rnn.num_units |
| input_size = self._rnn.input_size |
| |
| np_dtype = np.float32 if self._dtype == dtypes.float32 else np.float64 |
| inputs = np.random.randn(seq_length, batch_size, |
| input_size).astype(np_dtype) |
| input_h = np.random.randn(num_layers * dir_count, batch_size, |
| num_units).astype(np_dtype) |
| if self._rnn.rnn_mode == CUDNN_LSTM: |
| input_c = np.random.randn(num_layers * dir_count, batch_size, |
| num_units).astype(np_dtype) |
| initial_state = (input_h, input_c) |
| else: |
| initial_state = (input_h,) |
| return inputs, initial_state |
| |
| def ZeroState(self, batch_size): |
| num_layers = self._rnn.num_layers |
| dir_count = self._rnn.num_dirs |
| num_units = self._rnn.num_units |
| |
| np_dtype = np.float32 if self._dtype == dtypes.float32 else np.float64 |
| input_h = np.zeros((num_layers * dir_count, batch_size, |
| num_units)).astype(np_dtype) |
| if self._rnn.rnn_mode == CUDNN_LSTM: |
| input_c = np.zeros((num_layers * dir_count, batch_size, |
| num_units)).astype(np_dtype) |
| initial_state = (input_h, input_c) |
| else: |
| initial_state = (input_h,) |
| return initial_state |
| |
| def FProp(self, inputs_t, initial_state_t, training): |
| """Builds additional subgraph with given inputs and state. |
| |
| Args: |
| inputs_t: a tensor. |
| initial_state_t: a tensor. |
| training: boolean, true if training mode. |
| Returns: |
| A tensor of the forward pass output of the model. |
| """ |
| outputs, output_state = self._rnn( |
| inputs_t, initial_state=initial_state_t, training=training) |
| return self._AddUp(outputs, output_state) |
| |
| def Feed(self, sess, inputs, initial_state=None, return_sum=True): |
| """Runs graph with given inputs and initial state.""" |
| batch_size = inputs.shape[1] |
| if initial_state is None: |
| initial_state = self.ZeroState(batch_size) |
| if return_sum: |
| return sess.run( |
| self.total_sum, |
| feed_dict={self.inputs: inputs, |
| self.initial_state: initial_state}) |
| else: |
| return sess.run( |
| [self.outputs, self.output_state], |
| feed_dict={self.inputs: inputs, |
| self.initial_state: initial_state}) |
| |
| |
| def _CreateCudnnCompatibleCanonicalRNN(rnn, inputs, is_bidi=False, scope=None): |
| mode = rnn.rnn_mode |
| num_units = rnn.num_units |
| num_layers = rnn.num_layers |
| |
| # To reuse cuDNN-trained models, must use cudnn compatible rnn cells. |
| if mode == CUDNN_LSTM: |
| single_cell = lambda: cudnn_rnn_ops.CudnnCompatibleLSTMCell(num_units) |
| elif mode == CUDNN_GRU: |
| single_cell = lambda: cudnn_rnn_ops.CudnnCompatibleGRUCell(num_units) |
| elif mode == CUDNN_RNN_TANH: |
| single_cell = (lambda: rnn_cell_impl.BasicRNNCell(num_units, math_ops.tanh)) |
| elif mode == CUDNN_RNN_RELU: |
| single_cell = ( |
| lambda: rnn_cell_impl.BasicRNNCell(num_units, gen_nn_ops.relu)) |
| else: |
| raise ValueError("%s is not supported!" % mode) |
| |
| if not is_bidi: |
| cell = rnn_cell_impl.MultiRNNCell( |
| [single_cell() for _ in range(num_layers)]) |
| return rnn_lib.dynamic_rnn( |
| cell, inputs, dtype=dtypes.float32, time_major=True, scope=scope) |
| else: |
| cells_fw = [single_cell() for _ in range(num_layers)] |
| cells_bw = [single_cell() for _ in range(num_layers)] |
| |
| (outputs, output_state_fw, |
| output_state_bw) = contrib_rnn_lib.stack_bidirectional_dynamic_rnn( |
| cells_fw, |
| cells_bw, |
| inputs, |
| dtype=dtypes.float32, |
| time_major=True, |
| scope=scope) |
| return outputs, (output_state_fw, output_state_bw) |
| |
| |
| class CudnnRNNTestBasic(test_util.TensorFlowTestCase): |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testLayerBasic(self): |
| num_layers = 4 |
| num_units = 2 |
| batch_size = 8 |
| direction = CUDNN_RNN_UNIDIRECTION |
| dir_count = 1 |
| |
| with vs.variable_scope("main"): |
| kernel_initializer = init_ops.constant_initializer(0.) |
| bias_initializer = init_ops.constant_initializer(0.) |
| inputs = random_ops.random_uniform([ |
| num_layers * dir_count, batch_size, num_units], dtype=dtypes.float32) |
| |
| lstm = cudnn_rnn.CudnnLSTM(num_layers, num_units, |
| direction=direction, |
| kernel_initializer=kernel_initializer, |
| bias_initializer=bias_initializer, |
| name="awesome_lstm") |
| |
| # Build the layer |
| outputs1, _ = lstm(inputs) |
| # Reuse the layer |
| outputs2, _ = lstm(inputs) |
| |
| total_sum1 = math_ops.reduce_sum(outputs1) |
| total_sum2 = math_ops.reduce_sum(outputs2) |
| |
| with vs.variable_scope("main", reuse=True): |
| lstm = cudnn_rnn.CudnnLSTM(num_layers, num_units, |
| direction=direction, |
| kernel_initializer=kernel_initializer, |
| bias_initializer=bias_initializer, |
| name="awesome_lstm") |
| |
| # Reuse the layer |
| outputs3, _ = lstm(inputs) |
| total_sum3 = math_ops.reduce_sum(outputs3) |
| |
| self.assertEqual(1, len(variables.trainable_variables())) |
| self.assertEqual(1, len(ops.get_collection(ops.GraphKeys.SAVEABLE_OBJECTS))) |
| self.assertEqual("main/awesome_lstm/opaque_kernel", |
| variables.trainable_variables()[0].op.name) |
| |
| with self.test_session(use_gpu=True) as sess: |
| sess.run(variables.global_variables_initializer()) |
| (total_sum1_v, total_sum2_v, total_sum3_v) = sess.run( |
| [total_sum1, total_sum2, total_sum3]) |
| self.assertEqual(0, total_sum1_v) |
| self.assertEqual(0, total_sum2_v) |
| self.assertEqual(0, total_sum3_v) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testOptimizersSupport(self): |
| for opt in ("adagrad", "adam", "rmsprop", "momentum", "sgd"): |
| self._TestOptimizerSupportHelper(opt) |
| |
| def _GetOptimizer(self, opt): |
| if opt == "adagrad": |
| return adagrad.AdagradOptimizer(learning_rate=1e-2) |
| elif opt == "adam": |
| return adam.AdamOptimizer(learning_rate=1e-2) |
| elif opt == "rmsprop": |
| return rmsprop.RMSPropOptimizer(learning_rate=1e-2) |
| elif opt == "momentum": |
| return momentum.MomentumOptimizer(learning_rate=1e-2, momentum=0.9) |
| elif opt == "sgd": |
| return gradient_descent.GradientDescentOptimizer(learning_rate=1e-2) |
| else: |
| raise ValueError("Unsupported optimizer: %s" % opt) |
| |
| def _TestOptimizerSupportHelper(self, opt): |
| num_layers = 4 |
| num_units = 2 |
| batch_size = 8 |
| direction = CUDNN_RNN_UNIDIRECTION |
| dir_count = 1 |
| |
| with ops.Graph().as_default() as g: |
| kernel_initializer = init_ops.constant_initializer(0.) |
| bias_initializer = init_ops.constant_initializer(0.) |
| inputs = random_ops.random_uniform([ |
| num_layers * dir_count, batch_size, num_units], dtype=dtypes.float32) |
| |
| lstm = cudnn_rnn.CudnnLSTM(num_layers, num_units, |
| direction=direction, |
| kernel_initializer=kernel_initializer, |
| bias_initializer=bias_initializer, |
| name="awesome_lstm") |
| outputs, _ = lstm(inputs) |
| loss = math_ops.reduce_sum(outputs) |
| optimizer = self._GetOptimizer(opt) |
| train_op = optimizer.minimize(loss) |
| |
| with self.test_session(use_gpu=True, graph=g) as sess: |
| sess.run(variables.global_variables_initializer()) |
| sess.run(train_op) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSaveableGraphDeviceAssignment(self): |
| num_layers = 4 |
| num_units = 2 |
| batch_size = 8 |
| direction = CUDNN_RNN_UNIDIRECTION |
| dir_count = 1 |
| |
| def DeviceFn(op): |
| if op.type in ("Variable", "VariableV2"): |
| return "/cpu:0" |
| else: |
| return "/gpu:0" |
| |
| with ops.Graph().as_default() as g: |
| with ops.device(DeviceFn): |
| with vs.variable_scope("main"): |
| kernel_initializer = init_ops.constant_initializer(3.14) |
| bias_initializer = init_ops.constant_initializer(1.59) |
| inputs = random_ops.random_uniform( |
| [num_layers * dir_count, batch_size, num_units], |
| dtype=dtypes.float32) |
| |
| lstm = cudnn_rnn.CudnnLSTM(num_layers, num_units, |
| direction=direction, |
| kernel_initializer=kernel_initializer, |
| bias_initializer=bias_initializer, |
| name="awesome_lstm") |
| outputs = lstm(inputs) |
| |
| # saver is created in the scope of DeviceFn. |
| saver = saver_lib.Saver() |
| |
| with self.test_session(use_gpu=True, graph=g) as sess: |
| save_path = os.path.join(self.get_temp_dir(), |
| "test-saveable-device-assignment") |
| sess.run(variables.global_variables_initializer()) |
| |
| saver.save(sess, save_path) |
| saver.restore(sess, save_path) |
| sess.run(outputs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testDifferentShapesEager(self): |
| # Checks that kernel caching does not cause sharing of temporary storage |
| # across different input shapes when executing eagerly. |
| with context.eager_mode(): |
| with ops.device("gpu:0"): |
| first_output, _ = cudnn_rnn.CudnnGRU(1, 100)( |
| array_ops.zeros([28, 100, 28])) |
| second_output, _ = cudnn_rnn.CudnnGRU(1, 100)( |
| array_ops.zeros([28, 100, 100])) |
| self.assertAllEqual([28, 100, 100], first_output.shape) |
| self.assertAllEqual([28, 100, 100], second_output.shape) |
| |
| def _LossFunc(): |
| first_output, _ = cudnn_rnn.CudnnGRU(1, 100)( |
| array_ops.zeros([28, 100, 28])) |
| second_output, _ = cudnn_rnn.CudnnGRU(1, 100)( |
| array_ops.zeros([28, 100, 100])) |
| return (math_ops.reduce_sum(first_output) + |
| math_ops.reduce_sum(second_output)) |
| |
| backprop.implicit_grad(_LossFunc)() |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testDifferentShapesGraph(self): |
| # Tests that a single kernel instance presented with multiple input shapes |
| # does not crash with graph execution. |
| with ops.device("gpu:0"): |
| layer = cudnn_rnn.CudnnGRU(1, 100) |
| layer(array_ops.zeros([28, 100, 100])) |
| |
| def _Cond(index, accumulation): |
| del accumulation # unused |
| return math_ops.less(index, 4) |
| |
| def _Body(index, accumulation): |
| layer_input = accumulation[:, :, 10 * (1 + index % 2):] |
| output, _ = layer(layer_input) |
| return index + 1, accumulation + output |
| |
| original_input = array_ops.zeros([28, 100, 100]) |
| _, accumulation = control_flow_ops.while_loop(_Cond, _Body, |
| [0, original_input]) |
| grad, = gradients.gradients( |
| math_ops.reduce_sum(accumulation), (original_input,)) |
| init_op = variables.global_variables_initializer() |
| with self.test_session() as sess: |
| sess.run(init_op) |
| accumulation_eval, grad_eval = sess.run((accumulation, grad)) |
| self.assertAllEqual([28, 100, 100], accumulation_eval.shape) |
| self.assertAllEqual([28, 100, 100], grad_eval.shape) |
| |
| |
| # TODO(jamesqin): Transform to parameterized test after it is included in the |
| # TF open source codebase. |
| class CudnnRNNTestSaveRestore(test_util.TensorFlowTestCase): |
| |
| def _CompareWeights(self, lhs, rhs): |
| self.assertEqual(len(lhs), len(rhs)) |
| for lw, rw in zip(lhs, rhs): |
| self.assertAllEqual(lw, rw) |
| |
| def _CompareBiases(self, lhs, rhs, rnn_mode, num_layers, direction): |
| self.assertEqual(len(lhs), len(rhs)) |
| if rnn_mode == CUDNN_LSTM: |
| num_params_per_layer = CUDNN_LSTM_PARAMS_PER_LAYER |
| elif rnn_mode == CUDNN_GRU: |
| num_params_per_layer = CUDNN_GRU_PARAMS_PER_LAYER |
| elif rnn_mode == CUDNN_RNN_TANH: |
| num_params_per_layer = CUDNN_RNN_TANH_PARAMS_PER_LAYER |
| else: |
| num_params_per_layer = CUDNN_RNN_RELU_PARAMS_PER_LAYER |
| num_dirs = 1 if direction == CUDNN_RNN_UNIDIRECTION else 2 |
| num_params_per_layer *= num_dirs |
| self.assertEqual(num_params_per_layer * num_layers, len(lhs)) |
| |
| for i in range(num_layers): |
| layer_lhs = lhs[i * num_params_per_layer: (i+1) * num_params_per_layer] |
| layer_rhs = rhs[i * num_params_per_layer: (i+1) * num_params_per_layer] |
| if direction == CUDNN_RNN_UNIDIRECTION: |
| self._CompareSingleLayerBiases(layer_lhs, layer_rhs) |
| else: |
| size = len(layer_lhs) |
| fw_lhs, bw_lhs = layer_lhs[:size//2], layer_lhs[size//2:] |
| fw_rhs, bw_rhs = layer_rhs[:size//2], layer_rhs[size//2:] |
| self._CompareSingleLayerBiases(fw_lhs, fw_rhs) |
| self._CompareSingleLayerBiases(bw_lhs, bw_rhs) |
| |
| def _CompareSingleLayerBiases(self, lhs, rhs): |
| self.assertEqual(len(lhs), len(rhs)) |
| |
| lf_lhs, rt_lhs = lhs[:len(lhs)//2], lhs[len(lhs)//2:] |
| lf_rhs, rt_rhs = rhs[:len(rhs)//2], rhs[len(rhs)//2:] |
| self.assertEqual(len(lf_lhs), len(rt_lhs)) |
| self.assertEqual(len(lf_rhs), len(rt_rhs)) |
| |
| sum_lhs, sum_rhs = [], [] |
| for lf, rt in zip(lf_lhs, rt_lhs): |
| sum_lhs.append(lf + rt) |
| for lf, rt in zip(lf_rhs, rt_rhs): |
| sum_rhs.append(lf + rt) |
| self.assertEqual(len(sum_lhs), len(sum_rhs)) |
| for lf, rt in zip(sum_lhs, sum_rhs): |
| self.assertAllEqual(lf, rt) |
| |
| def _TestSaveRestoreVariable(self, rnn_mode, direction, dtype): |
| input_size = 3 |
| num_layers = 2 |
| num_units = 7 |
| with ops.Graph().as_default() as g: |
| random_seed.set_random_seed(1234) |
| model = CudnnTestModel( |
| rnn_mode, |
| num_layers, |
| num_units, |
| input_size, |
| direction=direction, |
| dtype=dtype) |
| rnn = model.rnn |
| save_path = os.path.join(self.get_temp_dir(), |
| "save-restore-variable-test") |
| saver = saver_lib.Saver() |
| weights, biases = model.rnn.saveable._OpaqueParamsToCanonical() |
| opaque_params = rnn.trainable_variables[0] |
| # CudnnTestModel() creates CudnnOpaqueParamsSaveable that helps saver save |
| # Cudnn vars in canonical format. |
| reset_op = state_ops.assign( |
| opaque_params, |
| array_ops.zeros(array_ops.shape(opaque_params), dtype=dtype)) |
| # Passing graph explicitly, otherwise an old sess would be reused. |
| with self.test_session(use_gpu=True, graph=g) as sess: |
| sess.run(variables.global_variables_initializer()) |
| val = saver.save(sess, save_path) |
| self.assertEqual(save_path, val) |
| weights_v, biases_v = sess.run([weights, biases]) |
| |
| # Reset opaque param |
| sess.run(reset_op) |
| saver.restore(sess, save_path) |
| weights_v_restored, biases_v_restored = sess.run([weights, biases]) |
| |
| self._CompareWeights(weights_v, weights_v_restored) |
| self._CompareBiases(biases_v, biases_v_restored, rnn_mode, num_layers, |
| direction) |
| |
| def _TestSaveRestoreTwoVariables(self, rnn_mode, direction, dtype): |
| input_size = 3 |
| num_layers = 2 |
| num_units = 7 |
| with ops.Graph().as_default() as g: |
| random_seed.set_random_seed(1234) |
| with vs.variable_scope("m1"): |
| model1 = CudnnTestModel( |
| rnn_mode, |
| num_layers, |
| num_units, |
| input_size, |
| direction=direction, |
| dtype=dtype) |
| with vs.variable_scope("m2"): |
| model2 = CudnnTestModel( |
| rnn_mode, |
| num_layers, |
| num_units, |
| input_size, |
| direction=direction, |
| dtype=dtype) |
| opaque_params = (model1.rnn.trainable_variables[0], |
| model2.rnn.trainable_variables[0]) |
| weights1, biases1 = model1.rnn.saveable._OpaqueParamsToCanonical() |
| weights2, biases2 = model2.rnn.saveable._OpaqueParamsToCanonical() |
| reset_params = [ |
| state_ops.assign(params, |
| array_ops.zeros_like(params, dtype=dtype)) |
| for params in opaque_params |
| ] |
| reset_op = control_flow_ops.group(*reset_params) |
| save_path = os.path.join(self.get_temp_dir(), |
| "save-restore-variable-test2") |
| saver = saver_lib.Saver() |
| # Passing graph explicitly, otherwise an old sess would be reused. |
| with self.test_session(use_gpu=True, graph=g) as sess: |
| sess.run(variables.global_variables_initializer()) |
| val = saver.save(sess, save_path) |
| self.assertEqual(save_path, val) |
| |
| weights1_v, biases1_v = sess.run([weights1, biases1]) |
| weights2_v, biases2_v = sess.run([weights2, biases2]) |
| |
| sess.run(reset_op) |
| saver.restore(sess, save_path) |
| weights1_v_restored, biases1_v_restored = sess.run([weights1, biases1]) |
| weights2_v_restored, biases2_v_restored = sess.run([weights2, biases2]) |
| |
| self._CompareWeights(weights1_v, weights1_v_restored) |
| self._CompareWeights(weights2_v, weights2_v_restored) |
| self._CompareBiases(biases1_v, biases1_v_restored, rnn_mode, num_layers, |
| direction) |
| self._CompareBiases(biases2_v, biases2_v_restored, rnn_mode, num_layers, |
| direction) |
| |
| def _TestSaveRestoreOutput(self, rnn_mode, direction, dtype): |
| with ops.Graph().as_default() as g: |
| num_layers = 2 |
| num_units = 7 |
| input_size = 7 |
| seq_length = 8 |
| batch_size = 4 |
| model = CudnnTestModel( |
| rnn_mode, |
| num_layers, |
| num_units, |
| input_size, |
| direction=direction, |
| dtype=dtype, |
| training=False) |
| rnn = model.rnn |
| |
| save_path = os.path.join(self.get_temp_dir(), "save-restore-output-test") |
| saver = saver_lib.Saver() |
| |
| # Only one opaque var in a cudnn layer. |
| assert len(rnn.trainable_variables) == 1 |
| reset_params = state_ops.assign( |
| rnn.trainable_variables[0], |
| array_ops.zeros( |
| array_ops.shape(rnn.trainable_variables[0]), dtype=dtype)) |
| |
| # Passing graph explicitly, otherwise an old sess would be reused. |
| with self.test_session(use_gpu=True, graph=g) as sess: |
| sess.run(variables.global_variables_initializer()) |
| inputs, initial_state = model.SynthesizeInput(seq_length, batch_size) |
| total_sum_v = model.Feed(sess, inputs, initial_state) |
| val = saver.save(sess, save_path) |
| self.assertEqual(save_path, val) |
| |
| sess.run(reset_params) |
| saver.restore(sess, save_path) |
| total_sum_v_restored = model.Feed(sess, inputs, initial_state) |
| self.assertAllClose(total_sum_v, total_sum_v_restored, atol=1e-5) |
| |
| def _TestSaveRestoreHelper(self, rnn_mode): |
| directions = [CUDNN_RNN_UNIDIRECTION, CUDNN_RNN_BIDIRECTION] |
| dtype_list = [dtypes.float16, dtypes.float32, dtypes.float64] |
| for direction, dtype in itertools.product(directions, dtype_list): |
| self._TestSaveRestoreVariable(rnn_mode, direction, dtype) |
| self._TestSaveRestoreTwoVariables(rnn_mode, direction, dtype) |
| self._TestSaveRestoreOutput(rnn_mode, direction, dtype) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSaveRestoreRepeatedlyCreateCustomSaveable(self): |
| input_size = 3 |
| num_layers = 2 |
| num_units = 7 |
| with ops.Graph().as_default(): |
| random_seed.set_random_seed(1234) |
| model = CudnnTestModel( |
| CUDNN_LSTM, |
| num_layers, |
| num_units, |
| input_size, |
| direction=CUDNN_RNN_UNIDIRECTION, |
| dtype=dtypes.float32) |
| with self.assertRaisesRegexp(RuntimeError, |
| "Cudnn saveable already created"): |
| model.rnn._create_saveable() |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSaveRestoreLSTM(self): |
| self._TestSaveRestoreHelper(CUDNN_LSTM) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSaveRestoreGRU(self): |
| self._TestSaveRestoreHelper(CUDNN_GRU) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSaveRestoreRNNTanh(self): |
| self._TestSaveRestoreHelper(CUDNN_RNN_TANH) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSaveRestoreRNNRelu(self): |
| self._TestSaveRestoreHelper(CUDNN_RNN_RELU) |
| |
| |
| class CudnnRNNTestSaveRestoreCheckpointable(test_util.TensorFlowTestCase): |
| |
| def _VerifyCheckpoint( |
| self, checkpoint_path, compatible_cell_fn, cudnn_cell_fn, |
| num_layers, input_size, expected_variable_values, num_applications=3): |
| checkpoint_directory = self.get_temp_dir() |
| checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") |
| with ops.device("gpu:0"): |
| cudnn_layer = cudnn_cell_fn() |
| cudnn_checkpoint = checkpointable_utils.Checkpoint(cell=cudnn_layer) |
| status = cudnn_checkpoint.restore(checkpoint_path) |
| inputs = 3. * array_ops.ones([num_applications, num_layers, input_size], |
| dtype=dtypes.float32) |
| cudnn_output, _ = cudnn_layer(inputs) |
| status.run_restore_ops() |
| second_save_path = cudnn_checkpoint.save(checkpoint_prefix) |
| restore_layer = compatible_cell_fn() |
| restore_layer_checkpoint = checkpointable_utils.Checkpoint( |
| cell=restore_layer) |
| status = restore_layer_checkpoint.restore(second_save_path) |
| current_state = restore_layer.zero_state(1, dtypes.float32) |
| for _ in range(num_applications): |
| restore_layer_output, current_state = restore_layer( |
| inputs=3. * array_ops.ones([1, input_size]), |
| state=current_state) |
| status.run_restore_ops() |
| self.assertTrue(restore_layer.variables) |
| for variable, expected_value in zip( |
| restore_layer.variables, expected_variable_values): |
| self.assertAllClose(expected_value, self.evaluate(variable)) |
| self.assertAllClose(self.evaluate(restore_layer_output), |
| self.evaluate(cudnn_output)[-1, -1:, ...]) |
| |
| def _CheckpointableSingleCellUnidirectionalTestTemplate( |
| self, single_cell_fn, cudnn_cell_fn): |
| # Single-layer cuDNN cells with object-based checkpointing should be |
| # checkpoint compatible with either single CudnnCompatible cells or |
| # MultiRnnCells with one cell. |
| input_size = 3 |
| save_cell_layer = single_cell_fn() |
| save_cell_layer( |
| inputs=array_ops.ones([1, input_size]), |
| state=save_cell_layer.zero_state(1, dtypes.float32)) |
| self.assertTrue(save_cell_layer.variables) |
| expected_values = [] |
| np.random.seed(10) |
| for variable in save_cell_layer.variables: |
| value = np.random.normal(size=variable.shape) |
| expected_values.append(value) |
| self.evaluate(variable.assign(value)) |
| save_checkpoint = checkpointable_utils.Checkpoint(cell=save_cell_layer) |
| checkpoint_directory = self.get_temp_dir() |
| checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") |
| first_save_path = save_checkpoint.save(checkpoint_prefix) |
| self._VerifyCheckpoint( |
| checkpoint_path=first_save_path, |
| compatible_cell_fn= |
| lambda: rnn_cell_impl.MultiRNNCell([single_cell_fn()]), |
| cudnn_cell_fn=cudnn_cell_fn, |
| num_layers=1, |
| expected_variable_values=expected_values, |
| input_size=input_size) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| @test_util.run_in_graph_and_eager_modes |
| def testLSTMCheckpointableSingleLayer(self): |
| num_units = 2 |
| direction = CUDNN_RNN_UNIDIRECTION |
| self._CheckpointableSingleCellUnidirectionalTestTemplate( |
| single_cell_fn=functools.partial( |
| cudnn_rnn_ops.CudnnCompatibleLSTMCell, num_units=num_units), |
| cudnn_cell_fn=functools.partial( |
| cudnn_rnn.CudnnLSTM, num_layers=1, num_units=num_units, |
| direction=direction, name="awesome_lstm")) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| @test_util.run_in_graph_and_eager_modes |
| def testGRUCheckpointableSingleLayer(self): |
| num_units = 2 |
| direction = CUDNN_RNN_UNIDIRECTION |
| with self.assertRaises(NotImplementedError): |
| # TODO(allenl): Implement object-based saving for GRUs and other cells. |
| self._CheckpointableSingleCellUnidirectionalTestTemplate( |
| single_cell_fn=functools.partial( |
| cudnn_rnn_ops.CudnnCompatibleGRUCell, num_units=num_units), |
| cudnn_cell_fn=functools.partial( |
| cudnn_rnn.CudnnGRU, num_layers=1, num_units=num_units, |
| direction=direction, name="awesome_gru")) |
| |
| def _CheckpointableMultiLayerTestTemplate( |
| self, single_cell_fn, cudnn_cell_fn, num_layers): |
| |
| def _MultiCellFn(): |
| return rnn_cell_impl.MultiRNNCell( |
| [single_cell_fn() for _ in range(num_layers)]) |
| input_size = 3 |
| save_graph = ops.Graph() |
| with save_graph.as_default(), self.session(graph=save_graph): |
| save_layer = _MultiCellFn() |
| save_layer(inputs=array_ops.ones([1, input_size]), |
| state=save_layer.zero_state(1, dtypes.float32)) |
| self.assertTrue(save_layer.variables) |
| expected_values = [] |
| np.random.seed(10) |
| for variable in save_layer.variables: |
| value = np.random.normal(size=variable.shape) |
| expected_values.append(value) |
| self.evaluate(variable.assign(value)) |
| save_checkpoint = checkpointable_utils.Checkpoint(cell=save_layer) |
| checkpoint_directory = self.get_temp_dir() |
| checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") |
| first_save_path = save_checkpoint.save(checkpoint_prefix) |
| self._VerifyCheckpoint( |
| checkpoint_path=first_save_path, |
| compatible_cell_fn=_MultiCellFn, cudnn_cell_fn=cudnn_cell_fn, |
| num_layers=num_layers, |
| expected_variable_values=expected_values, |
| input_size=input_size) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| @test_util.run_in_graph_and_eager_modes |
| def testCudnnCompatibleLSTMCheckpointablMultiLayer(self): |
| num_units = 2 |
| num_layers = 3 |
| direction = CUDNN_RNN_UNIDIRECTION |
| self._CheckpointableMultiLayerTestTemplate( |
| single_cell_fn=functools.partial( |
| cudnn_rnn_ops.CudnnCompatibleLSTMCell, num_units=num_units), |
| cudnn_cell_fn=functools.partial( |
| cudnn_rnn.CudnnLSTM, num_layers=num_layers, num_units=num_units, |
| direction=direction, name="awesome_lstm"), |
| num_layers=num_layers) |
| |
| |
| # TODO(jamesqin): Transform to parameterized test after it is included in the |
| # TF open source codebase. |
| class CudnnRNNTestCompatibleRNNCells(test_util.TensorFlowTestCase): |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testCudnnCompatibleLSTM(self): |
| self._TestCudnnCompatibleRnnCellsHelper(CUDNN_LSTM) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testCudnnCompatibleGRU(self): |
| self._TestCudnnCompatibleRnnCellsHelper(CUDNN_GRU) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testCudnnCompatibleRNNTanh(self): |
| self._TestCudnnCompatibleRnnCellsHelper(CUDNN_RNN_TANH) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testCudnnCompatibleRNNRelu(self): |
| self._TestCudnnCompatibleRnnCellsHelper(CUDNN_RNN_RELU) |
| |
| def _TestCudnnCompatibleRnnCellsHelper(self, rnn_mode): |
| configs = [ |
| { |
| "num_layers": 1, |
| "seq_length": 3, |
| "num_units": 4, |
| "input_size": 5, |
| "batch_size": 6, |
| }, |
| { |
| "num_layers": 2, |
| "seq_length": 8, |
| "num_units": 4, |
| "input_size": 8, |
| "batch_size": 16, |
| }, |
| { |
| "num_layers": 2, |
| "seq_length": 3, |
| "num_units": 4, |
| "input_size": 5, |
| "batch_size": 6, |
| }, |
| { |
| "num_layers": 1, |
| "seq_length": 2, |
| "num_units": 2, |
| "input_size": 4, |
| "batch_size": 1, |
| }, |
| ] |
| directions = [CUDNN_RNN_UNIDIRECTION, CUDNN_RNN_BIDIRECTION] |
| for cfg, direction in zip(configs, directions): |
| self._TestCudnnCompatibleRnnCells(cfg["num_layers"], cfg["seq_length"], |
| cfg["num_units"], cfg["input_size"], |
| cfg["batch_size"], rnn_mode, direction) |
| |
| def _TestCudnnCompatibleRnnCells(self, num_layers, seq_length, num_units, |
| input_size, batch_size, rnn_mode, direction): |
| dtype = dtypes.float32 |
| # Train graph |
| with ops.Graph().as_default() as g: |
| model = CudnnTestModel( |
| rnn_mode, |
| num_layers, |
| num_units, |
| input_size, |
| direction=direction, |
| dtype=dtype, |
| training=True) |
| target_output = array_ops.placeholder(dtype=dtype) |
| loss_op = losses.log_loss( |
| labels=target_output, predictions=model.total_sum) |
| optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=1e-2) |
| train_op = optimizer.minimize(loss_op) |
| |
| saver = saver_lib.Saver() |
| |
| # Train Cudnn model |
| seed = 0 |
| with self.test_session(use_gpu=True, graph=g) as sess: |
| sess.run(variables.global_variables_initializer()) |
| # Train 128 steps |
| num_steps = 128 |
| for _ in range(num_steps): |
| inputs, _ = model.SynthesizeInput(seq_length, batch_size, seed) |
| targets = np.random.rand() |
| sess.run( |
| train_op, |
| feed_dict={ |
| model.inputs: inputs, |
| model.initial_state: model.ZeroState(batch_size), |
| target_output: targets |
| }) |
| seed += 1 |
| |
| save_path = os.path.join(self.get_temp_dir(), |
| ("cudnn-rnn-%s-test" % rnn_mode)) |
| save_v = saver.save(sess, save_path) |
| self.assertEqual(save_path, save_v) |
| |
| # Cudnn inference graph |
| with ops.Graph().as_default() as g: |
| model = CudnnTestModel( |
| rnn_mode, |
| num_layers, |
| num_units, |
| input_size, |
| direction=direction, |
| dtype=dtype, |
| training=False) |
| rnn = model.rnn |
| saver = saver_lib.Saver() |
| |
| inference_input = np.random.rand(seq_length, batch_size, |
| input_size).astype(np.float32) |
| with self.test_session(use_gpu=True, graph=g) as sess: |
| sess.run(variables.global_variables_initializer()) |
| saver.restore(sess, save_path) |
| |
| # Cudnn inference |
| cudnn_outputs_v, cudnn_output_states_v = model.Feed( |
| sess, inference_input, return_sum=False) |
| |
| # Canonical RNN inference graph |
| with ops.Graph().as_default() as g: |
| cell_inputs = array_ops.placeholder( |
| dtype, shape=[seq_length, batch_size, input_size]) |
| if direction == CUDNN_RNN_UNIDIRECTION: |
| # outputs is one tensor, states are num_layer tuples, each 2 tensors |
| (outputs, states) = _CreateCudnnCompatibleCanonicalRNN(rnn, cell_inputs) |
| if rnn_mode == CUDNN_LSTM: |
| output_h = array_ops.stack([s.h for s in states]) |
| output_c = array_ops.stack([s.c for s in states]) |
| else: |
| output_state = array_ops.stack([s for s in states]) |
| else: |
| # outputs is one tensor. |
| # states is a tuple of 2 tuples: |
| # each sub tuple is num_layer tuples, each with 2 tensors. |
| (outputs, states) = _CreateCudnnCompatibleCanonicalRNN( |
| rnn, cell_inputs, is_bidi=True) |
| output_state_fw, output_state_bw = states |
| if rnn_mode == CUDNN_LSTM: |
| output_h, output_c = [], [] |
| for s_fw, s_bw in zip(output_state_fw, output_state_bw): |
| output_h.append(array_ops.stack([s_fw.h, s_bw.h])) |
| output_c.append(array_ops.stack([s_fw.c, s_bw.c])) |
| output_h = array_ops.concat(output_h, axis=0) |
| output_c = array_ops.concat(output_c, axis=0) |
| else: |
| output_state = [] |
| for s_fw, s_bw in zip(output_state_fw, output_state_bw): |
| output_state.append(array_ops.stack([s_fw, s_bw])) |
| output_state = array_ops.concat(output_state, axis=0) |
| saver = saver_lib.Saver() |
| |
| with self.test_session(use_gpu=True, graph=g) as sess: |
| saver.restore(sess, save_path) |
| |
| # BlockCell inference |
| if rnn_mode == CUDNN_LSTM: |
| outputs_v, output_h_v, output_c_v = sess.run( |
| [outputs, output_h, output_c], |
| feed_dict={cell_inputs: inference_input}) |
| self.assertAllClose(cudnn_outputs_v, outputs_v) |
| cudnn_output_h_v, cudnn_output_c_v = cudnn_output_states_v |
| self.assertAllClose(cudnn_output_h_v, output_h_v) |
| self.assertAllClose(cudnn_output_c_v, output_c_v) |
| else: |
| outputs_v, output_state_v = sess.run( |
| [outputs, output_state], |
| feed_dict={cell_inputs: inference_input}) |
| self.assertAllClose(cudnn_outputs_v, outputs_v, atol=2e-5, rtol=2e-5) |
| (cudnn_output_h_v,) = cudnn_output_states_v |
| self.assertAllClose(cudnn_output_h_v, output_state_v, atol=2e-5, |
| rtol=2e-5) |
| |
| |
| class CudnnRNNTestParamsSize(test_util.TensorFlowTestCase): |
| |
| def _TestOpaqueParamsSize(self, rnn_mode, num_layers, num_units, input_size, |
| dtype, direction): |
| logging.info("Testing one lstm param size with config: %s", locals()) |
| model = CudnnTestModel( |
| rnn_mode, |
| num_layers, |
| num_units, |
| input_size, |
| dtype=dtype, |
| direction=direction) |
| rnn = model.rnn |
| |
| # Min param size estimate = sum(weights.size) + sum(biases.size) |
| min_params_size = ( |
| np.sum(map(np.prod, rnn.canonical_weight_shapes)) + |
| np.sum([sp[0] for sp in rnn.canonical_bias_shapes])) |
| |
| opaque_params = rnn.trainable_variables[0] |
| with self.test_session(use_gpu=True, graph=ops.get_default_graph()): |
| variables.global_variables_initializer().run() |
| opaque_params_size_v = opaque_params.eval().size |
| self.assertLessEqual(min_params_size, opaque_params_size_v) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testOpaqueParamsSize(self): |
| test_configs = [ |
| [4, 200, 200], |
| [4, 200, 300], |
| [4, 200, 100], |
| [1, 100, 200], |
| [2, 200, 100], |
| [3, 200, 400], |
| ] |
| directions = [CUDNN_RNN_UNIDIRECTION, CUDNN_RNN_BIDIRECTION] |
| dtype_list = [dtypes.float16, dtypes.float32, dtypes.float64] |
| rnns = [CUDNN_LSTM, CUDNN_GRU, CUDNN_RNN_RELU, CUDNN_RNN_TANH] |
| for (rnn, config, dtype, direction) in itertools.product( |
| rnns, test_configs, dtype_list, directions): |
| num_layers, num_units, input_size = config |
| with ops.Graph().as_default(): |
| self._TestOpaqueParamsSize(rnn, num_layers, num_units, input_size, |
| dtype, direction) |
| |
| |
| class CudnnRNNTestTraining(test_util.TensorFlowTestCase): |
| |
| def setUp(self): |
| super(CudnnRNNTestTraining, self).setUp() |
| self._reset_rnd_gen_state = os.environ.get("TF_CUDNN_RESET_RND_GEN_STATE", |
| str(False)) |
| self._rnn_use_v2 = os.environ.get("TF_CUDNN_RNN_USE_V2", "0") |
| |
| def tearDown(self): |
| super(CudnnRNNTestTraining, self).tearDown() |
| os.environ["TF_CUDNN_RESET_RND_GEN_STATE"] = self._reset_rnd_gen_state |
| os.environ["TF_CUDNN_RNN_USE_V2"] = self._rnn_use_v2 |
| |
| def _ComputeNumericGrad(self, sess, y, x, delta=1e-4, step=1): |
| """Compute the numeric gradient of y wrt to x. |
| |
| Args: |
| sess: The TF session constructed with a graph containing x and y. |
| y: A scalar TF Tensor in the graph constructed in sess. |
| x: A TF Tensor in the graph constructed in sess. |
| delta: Gradient checker's small perturbation of x[i]. |
| step: Only compute numerical gradients for a subset of x values. |
| I.e. dy/dx[i] is computed if i % step == 0. |
| Returns: |
| A Tensor of the same shape and dtype as x. If x[i] is not chosen |
| to compute the numerical gradient dy/x[i], the corresponding |
| value is set to 0. |
| """ |
| |
| x_data = sess.run(x) |
| x_size = x_data.size |
| x_shape = x_data.shape |
| |
| numeric_grad = np.zeros(x_size, dtype=x_data.dtype) |
| |
| for i in range(0, x_size, step): |
| x_pos = x_data.copy() |
| if x_size == 1: |
| x_pos += delta |
| else: |
| x_pos.flat[i] += delta |
| y_pos_feed_dict = dict([(x.name, x_pos)]) |
| y_pos = sess.run(y, feed_dict=y_pos_feed_dict) |
| |
| x_neg = x_data.copy() |
| if x_size == 1: |
| x_neg -= delta |
| else: |
| x_neg.flat[i] -= delta |
| y_neg_feed_dict = dict([(x.name, x_neg)]) |
| y_neg = sess.run(y, feed_dict=y_neg_feed_dict) |
| numeric_grad[i] = (y_pos - y_neg) / (2 * delta) |
| return numeric_grad.reshape(x_shape) |
| |
| def _GetShape(self, sess, inputs): |
| if not isinstance(inputs, collections.Iterable): |
| return sess.run(array_ops.shape(inputs)) |
| else: |
| return sess.run([array_ops.shape(x) for x in inputs]) |
| |
| def _GradientCheckFp16(self, sess, y, xs, num_samples, |
| tolerance=1e-6, delta=1e-4): |
| """Gradient check for Fp16. |
| |
| Fp16 numerical gradients end up being zeros. Use a new way to check |
| gradients: |
| |
| Given multi-variant function: |
| y = f(x1, x2, ... xn) |
| delta_y = f(x1 + delta_x1, x2+delta_x2, ..., xn+delta_xn) - |
| f(x1, x2, ..., xn) |
| = f'(x1) * delta_x1 + f'(x2) * delta_x2 + .. + f'(xn) * delta_xn |
| where: |
| delta_xi are very small disturbance. |
| f'(xi) is the gradient of y w.r.t xi. |
| |
| The gradient check verifies the expected delta_y calculated by the above |
| equation is close to the actual delta_y. |
| Args: |
| sess: tf.Session object. |
| y: output tensor. |
| xs: a tensor or a list of input tensors. |
| num_samples: number of test samples to run. |
| tolerance: error tolerance. |
| delta: the order of magnititued of input disturbance to apply to calculate |
| the output change w.r.t inputs. |
| """ |
| sym_grads = self._ComputeSymGrads(sess, y, xs) |
| xs_shapes = self._GetShape(sess, xs) |
| |
| x_vals = [sess.run(x) for x in xs] |
| for _ in range(num_samples): |
| delta_xs = [delta * np.random.rand(*shape.tolist()) |
| for shape in xs_shapes] |
| |
| feed_dict = {} |
| for x, x_val, delta_x in zip(xs, x_vals, delta_xs): |
| feed_dict[x] = x_val + delta_x |
| actual_delta_y = (float(sess.run(y, feed_dict=feed_dict)) - |
| float(sess.run(y))) |
| |
| expected_delta_y = 0. |
| for sym_grad, delta_x in zip(sym_grads, delta_xs): |
| expected_delta_y += np.dot( |
| sym_grad.astype(np.float32).flatten(), |
| delta_x.astype(np.float32).flatten()) |
| self.assertAllClose(expected_delta_y, actual_delta_y, |
| atol=tolerance, rtol=tolerance) |
| |
| def _GradientCheck(self, sess, y, xs, tolerance=1e-6, delta=1e-4): |
| sym_grads = self._ComputeSymGrads(sess, y, xs) |
| |
| num_grads = [self._ComputeNumericGrad(sess, y, x, delta) for x in xs] |
| self.assertEqual(len(sym_grads), len(num_grads)) |
| for sym, num in zip(sym_grads, num_grads): |
| self.assertFalse(np.any(np.isnan(sym))) |
| self.assertFalse(np.any(np.isnan(num))) |
| self.assertAllClose(sym, num, atol=tolerance, rtol=tolerance) |
| |
| def _ComputeSymGrads(self, sess, y, xs): |
| sym_grads_t = gradients.gradients(y, xs) |
| return sess.run(sym_grads_t) |
| |
| def _TestOneSimpleTraining(self, rnn_mode, num_layers, num_units, input_size, |
| batch_size, seq_length, dir_count, dropout, dtype, |
| use_v2, delta, tolerance): |
| # Gradient checking runs two forward ops with almost the same input. Need to |
| # make sure the drop patterns across the two runs are the same. |
| logging.info("Training test with config: %s", locals()) |
| os.environ["TF_CUDNN_RESET_RND_GEN_STATE"] = str(True) |
| |
| np.random.seed(1234) |
| random_seed.set_random_seed(5678) |
| has_input_c = (rnn_mode == CUDNN_LSTM) |
| direction = (CUDNN_RNN_UNIDIRECTION |
| if dir_count == 1 else CUDNN_RNN_BIDIRECTION) |
| if use_v2: |
| os.environ["TF_CUDNN_RNN_USE_V2"] = "1" |
| else: |
| os.environ["TF_CUDNN_RNN_USE_V2"] = "0" |
| model = CudnnTestModel( |
| rnn_mode, |
| num_layers, |
| num_units, |
| input_size, |
| direction=direction, |
| dropout=dropout, |
| dtype=dtype, |
| training=True, |
| bias_initializer=init_ops.random_normal_initializer( |
| mean=1., dtype=dtype)) |
| rnn = model.rnn |
| params = rnn.trainable_variables[0] |
| |
| inputs = variables.Variable( |
| random_ops.random_uniform( |
| [seq_length, batch_size, input_size], dtype=dtype), |
| dtype=dtype) |
| input_h = variables.Variable( |
| random_ops.random_uniform( |
| [num_layers * dir_count, batch_size, num_units], dtype=dtype), |
| dtype=dtype) |
| if has_input_c: |
| input_c = variables.Variable( |
| random_ops.random_uniform( |
| [num_layers * dir_count, batch_size, num_units], dtype=dtype), |
| dtype=dtype) |
| initial_state = (input_h, input_c) |
| else: |
| initial_state = (input_h,) |
| total_sum = model.FProp(inputs, initial_state, training=True) |
| |
| with self.test_session(use_gpu=True, graph=ops.get_default_graph()) as sess: |
| sess.run(variables.global_variables_initializer()) |
| all_inputs = [inputs, params] |
| for s in initial_state: |
| all_inputs.append(s) |
| if dtype == dtypes.float16: |
| self._GradientCheckFp16( |
| sess, total_sum, all_inputs, |
| num_samples=FLAGS.grad_check_num_samples, |
| tolerance=tolerance, delta=delta) |
| else: |
| for _ in range(FLAGS.grad_check_num_samples): |
| # Each time choose a different set of inputs. |
| sess.run(variables.global_variables_initializer()) |
| self._GradientCheck( |
| sess, total_sum, all_inputs, |
| tolerance=tolerance, delta=delta) |
| |
| def _TestSimpleTrainingHelper(self, rnn_mode, test_configs): |
| dropouts = [0, 0.5, 1.] |
| v2_options = [str(False), str(True)] |
| for config, dropout, use_v2 in itertools.product(test_configs, dropouts, |
| v2_options): |
| dtype = config.get("dtype", dtypes.float32) |
| delta = config.get("delta", 1e-4) |
| tolerance = config.get("tolerance", 1e-6) |
| dir_count = config.get("dir_count", 1) |
| shape = config["shape"] |
| with ops.Graph().as_default(): |
| self._TestOneSimpleTraining( |
| rnn_mode, shape["num_layers"], shape["num_units"], |
| shape["input_size"], shape["batch_size"], shape["seq_length"], |
| dir_count, dropout, dtype, use_v2, delta, tolerance) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingLSTMFp64(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float64, |
| "tolerance": 5e-6, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_LSTM, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingLSTMFp32(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float32, |
| "delta": 1e-4, |
| "tolerance": 9e-2, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_LSTM, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingLSTMFp16(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float16, |
| "delta": 1e-3, |
| "tolerance": 9e-2, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| { |
| "dtype": dtypes.float16, |
| "delta": 1e-2, |
| "tolerance": 9e-2, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 6, |
| "input_size": 8, |
| "batch_size": 6, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_LSTM, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingGRUFp64(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float64, |
| "tolerance": 5e-6, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| } |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_GRU, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingGRUFp32(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float32, |
| "delta": 1e-3, |
| "tolerance": 4e-3, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_GRU, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingGRUFp16(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float16, |
| "delta": 2e-3, |
| "tolerance": 6e-2, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_GRU, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingRNNTanhFp64(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float64, |
| "tolerance": 5e-6, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_RNN_TANH, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingRNNTanhFp32(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float32, |
| "delta": 1e-3, |
| "tolerance": 5e-3, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_RNN_TANH, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingRNNTanhFp16(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float16, |
| "delta": 1e-3, |
| "tolerance": 5e-2, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_RNN_TANH, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingRNNReluFp64(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float64, |
| "tolerance": 5e-6, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_RNN_RELU, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingRNNReluFp32(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float32, |
| "delta": 1e-4, |
| "tolerance": 3e-1, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_RNN_RELU, test_configs) |
| |
| @unittest.skipUnless(test.is_built_with_cuda(), |
| "Test only applicable when running on GPUs") |
| def testSimpleTrainingRNNReluFp16(self): |
| test_configs = [ |
| { |
| "dtype": dtypes.float16, |
| "delta": 1e-3, |
| "tolerance": 7e-2, |
| "shape": { |
| "num_layers": 2, |
| "num_units": 3, |
| "input_size": 4, |
| "batch_size": 3, |
| "seq_length": 4, |
| }, |
| }, |
| ] |
| self._TestSimpleTrainingHelper(CUDNN_RNN_RELU, test_configs) |
| |
| |
| if __name__ == "__main__": |
| argv0 = sys.argv[0] |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| "--grad_check_num_samples", |
| type=int, |
| default=5, |
| help="Number of samples to run for gradient check.") |
| FLAGS, unparsed = parser.parse_known_args() |
| sys.argv = [argv0] + unparsed |
| googletest.main() |
