diff options
author | Shanqing Cai <cais@google.com> | 2017-12-25 23:53:51 -0500 |
---|---|---|
committer | GitHub <noreply@github.com> | 2017-12-25 23:53:51 -0500 |
commit | b95724bf529c1b39cd590e43fca762edff58b0f6 (patch) | |
tree | 60d7c76bb82877195fb35f3ba1a50c05966e47d2 | |
parent | e1ded7fa7cfacaeea43a903e738dd3fe2baabc57 (diff) | |
parent | 729be24efbedab5ef0a37ae8c19611d0d2aa6dce (diff) |
Merge pull request #15632 from drpngx/branch_180053468
Branch 180053468
3 files changed, 81 insertions, 37 deletions
diff --git a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py index e65394cba0..4ddd75de60 100644 --- a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py +++ b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py @@ -314,6 +314,47 @@ class CudnnRNNTestBasic(TensorFlowTestCase): self.assertEqual(0, total_sum2_v) self.assertEqual(0, total_sum3_v) + 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) + # TODO(jamesqin): Transform to parameterized test after it is included in the # TF open source codebase. diff --git a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py index 37c61a71a3..3a21914041 100644 --- a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py +++ b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py @@ -450,17 +450,18 @@ class _CudnnRNN(base_layer.Layer): raise RuntimeError( "%s._canonical_to_opaque invoked before input shape is known" % type(self).__name__) - return cudnn_rnn_ops.cudnn_rnn_canonical_to_opaque_params( - rnn_mode=self._rnn_mode, - num_layers=self._num_layers, - num_units=self._num_units, - input_size=self._input_size, - weights=cu_weights, - biases=cu_biases, - input_mode=self._input_mode, - seed=self._seed, - dropout=self._dropout, - direction=self._direction) + with ops.device("/gpu:0"): + return cudnn_rnn_ops.cudnn_rnn_canonical_to_opaque_params( + rnn_mode=self._rnn_mode, + num_layers=self._num_layers, + num_units=self._num_units, + input_size=self._input_size, + weights=cu_weights, + biases=cu_biases, + input_mode=self._input_mode, + seed=self._seed, + dropout=self._dropout, + direction=self._direction) def _forward(self, inputs, h, c, opaque_params, training): output, output_h, output_c = cudnn_rnn_ops._cudnn_rnn( # pylint:disable=protected-access @@ -489,12 +490,12 @@ class _CudnnRNN(base_layer.Layer): if self._saveable is not None: raise RuntimeError("Cudnn saveable already created.") self._saveable = self._saveable_cls( # pylint:disable=not-callable - self.trainable_variables[0], - self.num_layers, - self.num_units, - self.input_size, - self.input_mode, - self.direction, + opaque_params=self.trainable_variables[0], + num_layers=self.num_layers, + num_units=self.num_units, + input_size=self.input_size, + input_mode=self.input_mode, + direction=self.direction, scope=vs.get_variable_scope(), name="%s_saveable" % self.trainable_variables[0].op.name) ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, self._saveable) diff --git a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py index dcd3d4732a..45a3674bed 100644 --- a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py +++ b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py @@ -72,7 +72,7 @@ class CudnnCompatibleLSTMCell(lstm_ops.LSTMBlockCell): def __init__(self, num_units, reuse=None): super(CudnnCompatibleLSTMCell, self).__init__( num_units, forget_bias=0, cell_clip=None, use_peephole=False, - reuse=reuse) + reuse=reuse, name="cudnn_compatible_lstm_cell") self._names.update({"scope": "cudnn_compatible_lstm_cell"}) @@ -303,16 +303,17 @@ class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): Returns: 2 list for weights and biases respectively. """ - weights, biases = gen_cudnn_rnn_ops.cudnn_rnn_params_to_canonical( - num_layers=self._num_layers, - num_units=self._num_units, - input_size=self._input_size, - params=self._variables, - num_params=self._num_params, - rnn_mode=self._rnn_mode, - input_mode=self._input_mode, - direction=self._direction) - return (weights, biases) + with ops.device("/gpu:0"): + weights, biases = gen_cudnn_rnn_ops.cudnn_rnn_params_to_canonical( + num_layers=self._num_layers, + num_units=self._num_units, + input_size=self._input_size, + params=self._variables, + num_params=self._num_params, + rnn_mode=self._rnn_mode, + input_mode=self._input_mode, + direction=self._direction) + return (weights, biases) def _CanonicalToOpaqueParams(self, cu_weights, cu_biases): """Converts from Cudnn canonical format to opaque params. @@ -323,15 +324,16 @@ class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): Returns: a single opaque tensor. """ - return gen_cudnn_rnn_ops.cudnn_rnn_canonical_to_params( - num_layers=self._num_layers, - num_units=self._num_units, - input_size=self._input_size, - weights=cu_weights, - biases=cu_biases, - rnn_mode=self._rnn_mode, - input_mode=self._input_mode, - direction=self._direction) + with ops.device("/gpu:0"): + return gen_cudnn_rnn_ops.cudnn_rnn_canonical_to_params( + num_layers=self._num_layers, + num_units=self._num_units, + input_size=self._input_size, + weights=cu_weights, + biases=cu_biases, + rnn_mode=self._rnn_mode, + input_mode=self._input_mode, + direction=self._direction) def _TransformCanonical(self, cu_weights, cu_biases): r"""Transform from Cudnn canonical to tf canonical. |