experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Make LSTMCell use Defuns to speed up static graph builds, add compiled flag.

Browse Source 
Change: 145703555
This commit is contained in:
Eugene Brevdo 2017-01-26 12:06:38 -08:00 committed by TensorFlower Gardener
parent c76c5c4968
commit 13a1a9a71c
11 changed files with 388 additions and 158 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
configure vendored
View File

@ -168,10 +168,10 @@ done
if [ "$TF_ENABLE_XLA" == "1" ]; then
# Update Bazel build configuration.
perl -pi -e "s,WITH_XLA_SUPPORT = (False|True),WITH_XLA_SUPPORT = True,s" tensorflow/core/platform/default/build_config.bzl
sed -i -e "s/WITH_XLA_SUPPORT = (False|True)/WITH_XLA_SUPPORT = True/" tensorflow/core/platform/default/build_config_root.bzl
else
# Update Bazel build configuration.
perl -pi -e "s,WITH_XLA_SUPPORT = (False|True),WITH_XLA_SUPPORT = False,s" tensorflow/core/platform/default/build_config.bzl
sed -i -e "s/WITH_XLA_SUPPORT = (False|True)/WITH_XLA_SUPPORT = False/" tensorflow/core/platform/default/build_config_root.bzl
fi

2
tensorflow/contrib/rnn/BUILD
View File

@ -71,6 +71,7 @@ cuda_py_tests(
"//tensorflow/python:variable_scope",
"//tensorflow/python:variables",
],
xla_enabled = True,
)
cuda_py_tests(
@ -91,6 +92,7 @@ cuda_py_tests(
"//tensorflow/python:variable_scope",
"//tensorflow/python:variables",
],
xla_enabled = True,
)
cuda_py_tests(

165
tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py
View File

@ -19,6 +19,7 @@ from __future__ import division
from __future__ import print_function
import functools
import itertools
import sys
# TODO: #6568 Remove this hack that makes dlopen() not crash.
@ -33,9 +34,14 @@ import numpy as np
from tensorflow.contrib.rnn.python.ops import core_rnn_cell_impl
from tensorflow.contrib.rnn.python.ops.core_rnn_cell_impl import _linear as linear
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
@ -43,10 +49,41 @@ from tensorflow.python.ops import rnn
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.platform import test
from tensorflow.python.util import nest
# pylint: enable=protected-access
def _CreateMultiLSTMCellOps(batch_size, num_units, input_depth,
num_layers, max_time, compiled):
with variable_scope.variable_scope(
"root",
initializer=init_ops.random_uniform_initializer(-0.1, 0.1, seed=2)):
inputs = random_ops.random_uniform(
(max_time, batch_size, input_depth), seed=1)
rnn_cell = core_rnn_cell_impl.MultiRNNCell(
[core_rnn_cell_impl.LSTMCell(num_units, compiled=compiled)
for _ in range(num_layers)])
initial_state = rnn_cell.zero_state(
batch_size=batch_size, dtype=dtypes.float32)
outputs, final_state = rnn.dynamic_rnn(
cell=rnn_cell, inputs=inputs, initial_state=initial_state,
time_major=True)
flat_final_state = nest.flatten(final_state)
trainable_variables = variables_lib.trainable_variables()
outputs_grad = gradients_impl.gradients(
[outputs],
trainable_variables + [inputs] + nest.flatten(initial_state))
final_state_grad = gradients_impl.gradients(
flat_final_state,
trainable_variables + [inputs] + nest.flatten(initial_state))
return {"outputs": outputs,
"final_state": flat_final_state,
"outputs_grad": outputs_grad,
"final_state_grad": final_state_grad}
class RNNCellTest(test.TestCase):
def testLinear(self):
@ -117,8 +154,8 @@ class RNNCellTest(test.TestCase):
x = array_ops.zeros([1, 2])
m = array_ops.zeros([1, 8])
g, out_m = core_rnn_cell_impl.MultiRNNCell(
[core_rnn_cell_impl.BasicLSTMCell(
2, state_is_tuple=False)] * 2,
[core_rnn_cell_impl.BasicLSTMCell(2, state_is_tuple=False)
for _ in range(2)],
state_is_tuple=False)(x, m)
sess.run([variables_lib.global_variables_initializer()])
res = sess.run(
@ -165,7 +202,8 @@ class RNNCellTest(test.TestCase):
m0 = (array_ops.zeros([1, 2]),) * 2
m1 = (array_ops.zeros([1, 2]),) * 2
cell = core_rnn_cell_impl.MultiRNNCell(
[core_rnn_cell_impl.BasicLSTMCell(2)] * 2, state_is_tuple=True)
[core_rnn_cell_impl.BasicLSTMCell(2) for _ in range(2)],
state_is_tuple=True)
self.assertTrue(isinstance(cell.state_size, tuple))
self.assertTrue(
isinstance(cell.state_size[0], core_rnn_cell_impl.LSTMStateTuple))
@ -197,8 +235,8 @@ class RNNCellTest(test.TestCase):
m0 = array_ops.zeros([1, 4])
m1 = array_ops.zeros([1, 4])
cell = core_rnn_cell_impl.MultiRNNCell(
[core_rnn_cell_impl.BasicLSTMCell(
2, state_is_tuple=False)] * 2,
[core_rnn_cell_impl.BasicLSTMCell(2, state_is_tuple=False)
for _ in range(2)],
state_is_tuple=True)
g, (out_m0, out_m1) = cell(x, (m0, m1))
sess.run([variables_lib.global_variables_initializer()])
@ -407,7 +445,8 @@ class RNNCellTest(test.TestCase):
x = array_ops.zeros([1, 2])
m = array_ops.zeros([1, 4])
_, ml = core_rnn_cell_impl.MultiRNNCell(
[core_rnn_cell_impl.GRUCell(2)] * 2, state_is_tuple=False)(x, m)
[core_rnn_cell_impl.GRUCell(2) for _ in range(2)],
state_is_tuple=False)(x, m)
sess.run([variables_lib.global_variables_initializer()])
res = sess.run(ml, {
x.name: np.array([[1., 1.]]),
@ -416,6 +455,48 @@ class RNNCellTest(test.TestCase):
# The numbers in results were not calculated, this is just a smoke test.
self.assertAllClose(res, [[0.175991, 0.175991, 0.13248, 0.13248]])
def testMultiRNNCellWithLSTMCellAndXLA(self):
# TODO(b/34735319): Don't run this test if XLA is not available.
batch_size = 16
num_units = 32
input_depth = 12
num_layers = 2
max_time = 20
random_seed.set_random_seed(1234)
with self.test_session(graph=ops.Graph()) as sess:
xla_ops = _CreateMultiLSTMCellOps(
batch_size=batch_size, num_units=num_units,
input_depth=input_depth, num_layers=num_layers,
max_time=max_time,
compiled=True)
sess.run([variables_lib.global_variables_initializer()])
xla_results = sess.run(xla_ops)
random_seed.set_random_seed(1234)
with self.test_session(graph=ops.Graph()) as sess:
non_xla_ops = _CreateMultiLSTMCellOps(
batch_size=batch_size, num_units=num_units,
input_depth=input_depth, num_layers=num_layers,
max_time=max_time,
compiled=False)
sess.run([variables_lib.global_variables_initializer()])
non_xla_results = sess.run(non_xla_ops)
self.assertAllClose(non_xla_results["outputs"], xla_results["outputs"])
for xla_value, non_xla_value in zip(
xla_results["final_state"], non_xla_results["final_state"]):
self.assertAllClose(xla_value, non_xla_value)
for xla_g, non_xla_g in zip(
xla_results["outputs_grad"], non_xla_results["outputs_grad"]):
self.assertAllClose(xla_g, non_xla_g)
for xla_g, non_xla_g in zip(
xla_results["final_state_grad"], non_xla_results["final_state_grad"]):
self.assertAllClose(xla_g, non_xla_g)
def testMultiRNNCellWithStateTuple(self):
with self.test_session() as sess:
with variable_scope.variable_scope(
@ -427,11 +508,12 @@ class RNNCellTest(test.TestCase):
# Test incorrectness of state
with self.assertRaisesRegexp(ValueError, "Expected state .* a tuple"):
core_rnn_cell_impl.MultiRNNCell(
[core_rnn_cell_impl.GRUCell(2)] * 2,
[core_rnn_cell_impl.GRUCell(2) for _ in range(2)],
state_is_tuple=True)(x, m_bad)
_, ml = core_rnn_cell_impl.MultiRNNCell(
[core_rnn_cell_impl.GRUCell(2)] * 2, state_is_tuple=True)(x, m_good)
[core_rnn_cell_impl.GRUCell(2) for _ in range(2)],
state_is_tuple=True)(x, m_good)
sess.run([variables_lib.global_variables_initializer()])
res = sess.run(ml, {
@ -490,7 +572,7 @@ class SlimRNNCellTest(test.TestCase):
self.assertAllClose(res[1], res[3])
def basic_rnn_cell(inputs, state, num_units, scope=None):
def basic_rnn_cell(inputs, state, num_units, scope=None): # pylint: disable=invalid-name
if state is None:
if inputs is not None:
batch_size = inputs.get_shape()[0]
@ -512,5 +594,70 @@ def basic_rnn_cell(inputs, state, num_units, scope=None):
return output, output
class BenchmarkLSTMCellXLA(test.Benchmark):
def benchmarkDynamicRNNWithMultiLSTMCell(self):
num_layers = 3
max_time = 50
print("benchmarkDynamicRNNWithMultiLSTMCell")
print("\t" +
"\t".join(["inter_th", "intra_th",
"batch_size", "num_units", "input_depth", "device",
"compiled", "wall_time"]))
warmup_run = True
for (threads,
device,
num_units,
batch_size,
input_depth,
compiled) in itertools.product(
[{"inter": 0, "intra": 0}, {"inter": 1, "intra": 4}],
["cpu", "gpu"],
[32, 512],
[1, 32, 256],
[32, 512],
[False, True]):
if threads["inter"] != 0:
# We only care about testing inter/intra op limitations on
# CPU with small batch size, to mimic embedded devices.
if device != "cpu" or batch_size != 1:
continue
if device == "cpu" and batch_size > 32:
continue
random_seed.set_random_seed(1234)
config = config_pb2.ConfigProto(
inter_op_parallelism_threads=threads["inter"],
intra_op_parallelism_threads=threads["intra"],
allow_soft_placement=False)
with session.Session(config=config, graph=ops.Graph()) as sess:
with ops.device("/%s:0" % device):
ops_dict = _CreateMultiLSTMCellOps(
batch_size=batch_size, num_units=num_units,
input_depth=input_depth, num_layers=num_layers,
max_time=max_time,
compiled=compiled)
sess.run([variables_lib.global_variables_initializer()])
all_ops = nest.flatten(ops_dict.values())
all_ops_group = control_flow_ops.group(*all_ops)
name_suffix = (
"inter_th_%d_intra_th_%d_bs_%d_units_%d_inputdepth_%d"
"_device_%s_xla_%s" % (
threads["inter"], threads["intra"],
batch_size, num_units, input_depth, device, compiled))
if warmup_run:
self.run_op_benchmark(
sess, all_ops_group, min_iters=30, name="ignore_warmup")
warmup_run = False
benchmark_results = self.run_op_benchmark(
sess, all_ops_group, min_iters=30,
name="benchmarkDynamicRNNWithMultiLSTMCell_%s" % name_suffix)
print("\t" +
"\t".join(["%s" % x for x in [
threads["inter"], threads["intra"],
batch_size, num_units, input_depth, device, compiled,
benchmark_results["wall_time"]]]))
if __name__ == "__main__":
test.main()

103
tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py
View File

@ -154,6 +154,7 @@ class RNNTest(test.TestCase):
def setUp(self):
self._seed = 23489
np.random.seed(self._seed)
ops_lib.reset_default_graph()
def testInvalidSequenceLengthShape(self):
cell = Plus1RNNCell()
@ -583,7 +584,7 @@ class LSTMTest(test.TestCase):
(state_notuple_v,) = sess.run((state_notuple,),
feed_dict={inputs[0]: input_value})
state_tuple_v = sess.run(state_tuple, feed_dict={inputs[0]: input_value})
self.assertAllEqual(state_notuple_v, np.hstack(state_tuple_v))
self.assertAllClose(state_notuple_v, np.hstack(state_tuple_v))
def _testProjSharding(self, use_gpu):
num_units = 3
@ -806,7 +807,7 @@ class LSTMTest(test.TestCase):
self.assertEqual(len(outputs0_values), len(outputs2_values))
for o1, o2, o3 in zip(outputs0_values, outputs1_values, outputs2_values):
# Same weights used by both RNNs so outputs should be the same.
self.assertAllEqual(o1, o2)
self.assertAllClose(o1, o2)
# Different weights used so outputs should be different.
self.assertTrue(np.linalg.norm(o1 - o3) > 1e-6)
@ -844,7 +845,7 @@ class LSTMTest(test.TestCase):
outputs1_values = output_values[max_length:]
self.assertEqual(len(outputs0_values), len(outputs1_values))
for out0, out1 in zip(outputs0_values, outputs1_values):
self.assertAllEqual(out0, out1)
self.assertAllClose(out0, out1)
def testNoProjNoShardingSimpleStateSaver(self):
self._testNoProjNoShardingSimpleStateSaver(use_gpu=False)
@ -934,13 +935,13 @@ class LSTMTest(test.TestCase):
feed_dict={inputs[0]: input_value})
outputs_dynamic_v = sess.run(outputs_dynamic,
feed_dict={inputs[0]: input_value})
self.assertAllEqual(outputs_static_v, outputs_dynamic_v)
self.assertAllClose(outputs_static_v, outputs_dynamic_v)
state_static_v = sess.run(state_static,
feed_dict={inputs[0]: input_value})
state_dynamic_v = sess.run(state_dynamic,
feed_dict={inputs[0]: input_value})
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_dynamic_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_dynamic_v))
def testDynamicRNNWithNestedTupleStates(self):
num_units = 3
@ -1003,13 +1004,13 @@ class LSTMTest(test.TestCase):
feed_dict={inputs[0]: input_value})
outputs_dynamic_v = sess.run(outputs_dynamic,
feed_dict={inputs[0]: input_value})
self.assertAllEqual(outputs_static_v, outputs_dynamic_v)
self.assertAllClose(outputs_static_v, outputs_dynamic_v)
state_static_v = sess.run(nest.flatten(state_static),
feed_dict={inputs[0]: input_value})
state_dynamic_v = sess.run(nest.flatten(state_dynamic),
feed_dict={inputs[0]: input_value})
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_dynamic_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_dynamic_v))
def _testDynamicEquivalentToStaticRNN(self, use_gpu, use_sequence_length):
time_steps = 8
@ -1038,7 +1039,9 @@ class LSTMTest(test.TestCase):
use_peepholes=True,
initializer=initializer,
num_proj=num_proj,
state_is_tuple=False)
state_is_tuple=False,
# TODO(b/XXX): Defun name aliasing causes errors
compiled=False)
with variable_scope.variable_scope("dynamic_scope"):
outputs_static, state_static = core_rnn.static_rnn(
@ -1096,7 +1099,9 @@ class LSTMTest(test.TestCase):
use_peepholes=True,
initializer=initializer,
num_proj=num_proj,
state_is_tuple=False)
state_is_tuple=False,
# TODO(b/XXX): Defun name aliasing causes errors
compiled=False)
with variable_scope.variable_scope("dynamic_scope"):
outputs_dynamic, state_dynamic = rnn.dynamic_rnn(
@ -1150,10 +1155,10 @@ class LSTMTest(test.TestCase):
######### Step 3: Comparisons
self.assertEqual(len(values_static), len(values_dynamic))
for (value_static, value_dynamic) in zip(values_static, values_dynamic):
self.assertAllEqual(value_static, value_dynamic)
self.assertAllEqual(state_value_static, state_value_dynamic)
self.assertAllClose(value_static, value_dynamic)
self.assertAllClose(state_value_static, state_value_dynamic)
self.assertAllEqual(static_grad_values, dynamic_grad_values)
self.assertAllClose(static_grad_values, dynamic_grad_values)
self.assertEqual(
len(static_individual_grad_values), len(dynamic_individual_grad_values))
@ -1164,14 +1169,14 @@ class LSTMTest(test.TestCase):
for i, (a, b) in enumerate(
zip(static_individual_grad_values, dynamic_individual_grad_values)):
tf_logging.info("Comparing individual gradients iteration %d" % i)
self.assertAllEqual(a, b)
self.assertAllClose(a, b)
for i, (a, b) in enumerate(
zip(static_individual_var_grad_values,
dynamic_individual_var_grad_values)):
tf_logging.info("Comparing individual variable gradients iteration %d" %
i)
self.assertAllEqual(a, b)
self.assertAllClose(a, b)
def testDynamicEquivalentToStaticRNN(self):
self._testDynamicEquivalentToStaticRNN(
@ -1293,13 +1298,13 @@ class BidirectionalRNNTest(test.TestCase):
# Both sequences in batch are length=8. Check that the time=i
# forward output is equal to time=8-1-i backward output
for i in xrange(8):
self.assertEqual(out[i][0][0], out[8 - 1 - i][0][3])
self.assertEqual(out[i][0][1], out[8 - 1 - i][0][4])
self.assertEqual(out[i][0][2], out[8 - 1 - i][0][5])
self.assertAllClose(out[i][0][0], out[8 - 1 - i][0][3])
self.assertAllClose(out[i][0][1], out[8 - 1 - i][0][4])
self.assertAllClose(out[i][0][2], out[8 - 1 - i][0][5])
for i in xrange(8):
self.assertEqual(out[i][1][0], out[8 - 1 - i][1][3])
self.assertEqual(out[i][1][1], out[8 - 1 - i][1][4])
self.assertEqual(out[i][1][2], out[8 - 1 - i][1][5])
self.assertAllClose(out[i][1][0], out[8 - 1 - i][1][3])
self.assertAllClose(out[i][1][1], out[8 - 1 - i][1][4])
self.assertAllClose(out[i][1][2], out[8 - 1 - i][1][5])
# Via the reasoning above, the forward and backward final state should be
# exactly the same
self.assertAllClose(s_fw, s_bw)
@ -1399,27 +1404,27 @@ class BidirectionalRNNTest(test.TestCase):
# Check that the time=0 forward output is equal to time=1 backward output
if not use_time_major:
out = np.swapaxes(out, 0, 1)
self.assertEqual(out[0][0][0], out[1][0][3])
self.assertEqual(out[0][0][1], out[1][0][4])
self.assertEqual(out[0][0][2], out[1][0][5])
self.assertAllClose(out[0][0][0], out[1][0][3])
self.assertAllClose(out[0][0][1], out[1][0][4])
self.assertAllClose(out[0][0][2], out[1][0][5])
# Check that the time=1 forward output is equal to time=0 backward output
self.assertEqual(out[1][0][0], out[0][0][3])
self.assertEqual(out[1][0][1], out[0][0][4])
self.assertEqual(out[1][0][2], out[0][0][5])
self.assertAllClose(out[1][0][0], out[0][0][3])
self.assertAllClose(out[1][0][1], out[0][0][4])
self.assertAllClose(out[1][0][2], out[0][0][5])
# Second sequence in batch is length=3
# Check that the time=0 forward output is equal to time=2 backward output
self.assertEqual(out[0][1][0], out[2][1][3])
self.assertEqual(out[0][1][1], out[2][1][4])
self.assertEqual(out[0][1][2], out[2][1][5])
self.assertAllClose(out[0][1][0], out[2][1][3])
self.assertAllClose(out[0][1][1], out[2][1][4])
self.assertAllClose(out[0][1][2], out[2][1][5])
# Check that the time=1 forward output is equal to time=1 backward output
self.assertEqual(out[1][1][0], out[1][1][3])
self.assertEqual(out[1][1][1], out[1][1][4])
self.assertEqual(out[1][1][2], out[1][1][5])
self.assertAllClose(out[1][1][0], out[1][1][3])
self.assertAllClose(out[1][1][1], out[1][1][4])
self.assertAllClose(out[1][1][2], out[1][1][5])
# Check that the time=2 forward output is equal to time=0 backward output
self.assertEqual(out[2][1][0], out[0][1][3])
self.assertEqual(out[2][1][1], out[0][1][4])
self.assertEqual(out[2][1][2], out[0][1][5])
self.assertAllClose(out[2][1][0], out[0][1][3])
self.assertAllClose(out[2][1][1], out[0][1][4])
self.assertAllClose(out[2][1][2], out[0][1][5])
# Via the reasoning above, the forward and backward final state should be
# exactly the same
self.assertAllClose(s_fw, s_bw)
@ -1560,13 +1565,13 @@ class MultiDimensionalLSTMTest(test.TestCase):
outputs_sav_v = sess.run(outputs_sav,
feed_dict={inputs_using_dim[0]: input_value})
self.assertAllEqual(outputs_static_v, outputs_dynamic_v)
self.assertAllEqual(outputs_static_v, outputs_sav_v)
self.assertAllClose(outputs_static_v, outputs_dynamic_v)
self.assertAllClose(outputs_static_v, outputs_sav_v)
outputs_static_array = np.array(outputs_static_v)
outputs_static_array_double = np.concatenate(
(outputs_static_array, outputs_static_array), axis=2)
outputs_bid_array = np.array(outputs_bid_v)
self.assertAllEqual(outputs_static_array_double, outputs_bid_array)
self.assertAllClose(outputs_static_array_double, outputs_bid_array)
state_static_v = sess.run(state_static,
feed_dict={inputs[0]: input_value})
@ -1578,10 +1583,10 @@ class MultiDimensionalLSTMTest(test.TestCase):
feed_dict={inputs_using_dim[0]: input_value})
state_sav_v = sess.run(state_sav,
feed_dict={inputs_using_dim[0]: input_value})
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_dynamic_v))
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_sav_v))
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_bid_fw_v))
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_bid_bw_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_dynamic_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_sav_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_bid_fw_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_bid_bw_v))
class NestedLSTMTest(test.TestCase):
@ -1663,14 +1668,14 @@ class NestedLSTMTest(test.TestCase):
outputs_bid_v = sess.run(outputs_bid,
feed_dict={single_input_using_dim: input_value})
self.assertAllEqual(outputs_static_v,
self.assertAllClose(outputs_static_v,
np.transpose(outputs_dynamic_v, (1, 0, 2, 3)))
self.assertAllEqual(outputs_static_v, outputs_sav_v)
self.assertAllClose(outputs_static_v, outputs_sav_v)
outputs_static_array = np.array(outputs_static_v)
outputs_static_array_double = np.concatenate(
(outputs_static_array, outputs_static_array), axis=3)
outputs_bid_array = np.array(outputs_bid_v)
self.assertAllEqual(outputs_static_array_double, outputs_bid_array)
self.assertAllClose(outputs_static_array_double, outputs_bid_array)
state_dynamic_v = sess.run(state_dynamic,
feed_dict={single_input: input_value})
@ -1682,10 +1687,10 @@ class NestedLSTMTest(test.TestCase):
feed_dict={single_input_using_dim: input_value})
state_sav_v = sess.run(state_sav,
feed_dict={single_input_using_dim: input_value})
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_dynamic_v))
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_sav_v))
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_bid_fw_v))
self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_bid_bw_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_dynamic_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_sav_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_bid_fw_v))
self.assertAllClose(np.hstack(state_static_v), np.hstack(state_bid_bw_v))
class StateSaverRNNTest(test.TestCase):

111
tensorflow/contrib/rnn/python/ops/core_rnn_cell_impl.py
View File

@ -22,6 +22,7 @@ from __future__ import print_function
import collections
import math
from tensorflow.python.framework import function
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import clip_ops
@ -61,7 +62,7 @@ class BasicRNNCell(RNNCell):
"""Most basic RNN: output = new_state = act(W * input + U * state + B)."""
with vs.variable_scope(scope or "basic_rnn_cell"):
output = self._activation(
_linear([inputs, state], self._num_units, True, scope=scope))
_linear([inputs, state], self._num_units, True))
return output, output
@ -89,14 +90,13 @@ class GRUCell(RNNCell):
# We start with bias of 1.0 to not reset and not update.
r, u = array_ops.split(
value=_linear(
[inputs, state], 2 * self._num_units, True, 1.0, scope=scope),
[inputs, state], 2 * self._num_units, True, 1.0),
num_or_size_splits=2,
axis=1)
r, u = sigmoid(r), sigmoid(u)
with vs.variable_scope("candidate"):
c = self._activation(_linear([inputs, r * state],
self._num_units, True,
scope=scope))
self._num_units, True))
new_h = u * state + (1 - u) * c
return new_h, new_h
@ -176,7 +176,7 @@ class BasicLSTMCell(RNNCell):
c, h = state
else:
c, h = array_ops.split(value=state, num_or_size_splits=2, axis=1)
concat = _linear([inputs, h], 4 * self._num_units, True, scope=scope)
concat = _linear([inputs, h], 4 * self._num_units, True)
# i = input_gate, j = new_input, f = forget_gate, o = output_gate
i, j, f, o = array_ops.split(value=concat, num_or_size_splits=4, axis=1)
@ -192,6 +192,13 @@ class BasicLSTMCell(RNNCell):
return new_h, new_state
def _maybe_compile(fun, compiled):
if not compiled:
return fun
else:
return function.Defun(noinline=True, compiled=True)(fun)
class LSTMCell(RNNCell):
"""Long short-term memory unit (LSTM) recurrent network cell.
@ -219,7 +226,7 @@ class LSTMCell(RNNCell):
initializer=None, num_proj=None, proj_clip=None,
num_unit_shards=None, num_proj_shards=None,
forget_bias=1.0, state_is_tuple=True,
activation=tanh):
activation=tanh, compiled=False):
"""Initialize the parameters for an LSTM cell.
Args:
@ -246,6 +253,12 @@ class LSTMCell(RNNCell):
the `c_state` and `m_state`. If False, they are concatenated
along the column axis. This latter behavior will soon be deprecated.
activation: Activation function of the inner states.
compiled: Python boolean. If `True`, the core computation of the LSTM
cell is compiled via XLA. As of now, this provides speedups for
most GPU calculations, and on small batch CPU and embedded calculations.
Raises:
ValueError: if compiled=True and state_is_tuple=False (not supported).
"""
if not state_is_tuple:
logging.warn("%s: Using a concatenated state is slower and will soon be "
@ -257,6 +270,9 @@ class LSTMCell(RNNCell):
"%s: The num_unit_shards and proj_unit_shards parameters are "
"deprecated and will be removed in Jan 2017. "
"Use a variable scope with a partitioner instead.", self)
if not state_is_tuple and compiled:
raise ValueError(
"Combining state_is_tuple=False and compiled=True is not supported.")
self._num_units = num_units
self._use_peepholes = use_peepholes
@ -269,6 +285,7 @@ class LSTMCell(RNNCell):
self._forget_bias = forget_bias
self._state_is_tuple = state_is_tuple
self._activation = activation
self._compiled = compiled
if num_proj:
self._state_size = (
@ -317,16 +334,35 @@ class LSTMCell(RNNCell):
"""
num_proj = self._num_units if self._num_proj is None else self._num_proj
if self._state_is_tuple:
(c_prev, m_prev) = state
else:
c_prev = array_ops.slice(state, [0, 0], [-1, self._num_units])
m_prev = array_ops.slice(state, [0, self._num_units], [-1, num_proj])
def _kernel(k_inputs, state_p0, state_p1):
"""Internal kernel for the single step of LSTM.
dtype = inputs.dtype
input_size = inputs.get_shape().with_rank(2)[1]
Args:
k_inputs: Input Tensor.
state_p0: Either the state or the c component of the state.
state_p1: Either the state or the m component of the state.
Returns:
(m, c) or (m, concat([c, m])) depending on state_is_tuple.
Raises:
ValueError: see above docstring.
"""
k_inputs.set_shape(inputs.get_shape())
if self._state_is_tuple:
(c_prev, m_prev) = state_p0, state_p1
c_prev.set_shape(state[0].get_shape())
m_prev.set_shape(state[1].get_shape())
else:
k_state = state_p0
c_prev = array_ops.slice(k_state, [0, 0], [-1, self._num_units])
m_prev = array_ops.slice(k_state, [0, self._num_units], [-1, num_proj])
dtype = k_inputs.dtype
input_size = k_inputs.get_shape().with_rank(2)[1]
if input_size.value is None:
raise ValueError("Could not infer input size from inputs.get_shape()[-1]")
raise ValueError(
"Could not infer input size from inputs.get_shape()[-1]")
with vs.variable_scope(scope or "lstm_cell",
initializer=self._initializer) as unit_scope:
if self._num_unit_shards is not None:
@ -334,8 +370,9 @@ class LSTMCell(RNNCell):
partitioned_variables.fixed_size_partitioner(
self._num_unit_shards))
# i = input_gate, j = new_input, f = forget_gate, o = output_gate
lstm_matrix = _linear([inputs, m_prev], 4 * self._num_units, bias=True,
scope=scope)
lstm_matrix = _linear(
[k_inputs, m_prev], 4 * self._num_units, bias=True,
compiled=self._compiled)
i, j, f, o = array_ops.split(
value=lstm_matrix, num_or_size_splits=4, axis=1)
@ -350,8 +387,6 @@ class LSTMCell(RNNCell):
"w_i_diag", shape=[self._num_units], dtype=dtype)
w_o_diag = vs.get_variable(
"w_o_diag", shape=[self._num_units], dtype=dtype)
if self._use_peepholes:
c = (sigmoid(f + self._forget_bias + w_f_diag * c_prev) * c_prev +
sigmoid(i + w_i_diag * c_prev) * self._activation(j))
else:
@ -374,16 +409,36 @@ class LSTMCell(RNNCell):
proj_scope.set_partitioner(
partitioned_variables.fixed_size_partitioner(
self._num_proj_shards))
m = _linear(m, self._num_proj, bias=False, scope=scope)
m = _linear(m, self._num_proj, bias=False, compiled=self._compiled)
if self._proj_clip is not None:
# pylint: disable=invalid-unary-operand-type
m = clip_ops.clip_by_value(m, -self._proj_clip, self._proj_clip)
# pylint: enable=invalid-unary-operand-type
new_state = (LSTMStateTuple(c, m) if self._state_is_tuple else
array_ops.concat([c, m], 1))
return m, new_state
if self._state_is_tuple:
return m, c
else:
return m, array_ops.concat([c, m], 1)
compiled_kernel = _maybe_compile(_kernel, self._compiled)
if self._state_is_tuple:
batch_shape = (
inputs.get_shape()[:1].merge_with(
state[0].get_shape()[:1]).merge_with(
state[1].get_shape()[:1]))
emit_m, emit_c = compiled_kernel(inputs, state[0], state[1])
emit_c.set_shape(batch_shape.concatenate([state[0].get_shape()[1]]))
emit_m.set_shape(batch_shape.concatenate([state[1].get_shape()[1]]))
emit_state = LSTMStateTuple(emit_c, emit_m)
else:
batch_shape = inputs.get_shape()[:1].merge_with(state.get_shape()[:1])
emit_m, emit_state = compiled_kernel(inputs, state, state)
emit_m.set_shape(batch_shape.concatenate([num_proj]))
emit_state.set_shape(batch_shape.concatenate([state.get_shape()[1]]))
return emit_m, emit_state
class OutputProjectionWrapper(RNNCell):
@ -426,7 +481,7 @@ class OutputProjectionWrapper(RNNCell):
output, res_state = self._cell(inputs, state)
# Default scope: "OutputProjectionWrapper"
with vs.variable_scope(scope or "output_projection_wrapper"):
projected = _linear(output, self._output_size, True, scope=scope)
projected = _linear(output, self._output_size, True)
return projected, res_state
@ -468,7 +523,7 @@ class InputProjectionWrapper(RNNCell):
"""Run the input projection and then the cell."""
# Default scope: "InputProjectionWrapper"
with vs.variable_scope(scope or "input_projection_wrapper"):
projected = _linear(inputs, self._num_proj, True, scope=scope)
projected = _linear(inputs, self._num_proj, True)
return self._cell(projected, state)
@ -762,7 +817,7 @@ class _SlimRNNCell(RNNCell):
return output, state
def _linear(args, output_size, bias, bias_start=0.0, scope=None):
def _linear(args, output_size, bias, bias_start=0.0, compiled=False):
"""Linear map: sum_i(args[i] * W[i]), where W[i] is a variable.
Args:
@ -770,7 +825,7 @@ def _linear(args, output_size, bias, bias_start=0.0, scope=None):
output_size: int, second dimension of W[i].
bias: boolean, whether to add a bias term or not.
bias_start: starting value to initialize the bias; 0 by default.
scope: (optional) Variable scope to create parameters in.
compiled: boolean, _linear plays nicely with XLA if it is enabled.
Returns:
A 2D Tensor with shape [batch x output_size] equal to
@ -815,4 +870,8 @@ def _linear(args, output_size, bias, bias_start=0.0, scope=None):
"biases", [output_size],
dtype=dtype,
initializer=init_ops.constant_initializer(bias_start, dtype=dtype))
if compiled:
# TODO(b/34505635): Defuns don't play well with bias_add
return res + biases
else:
return nn_ops.bias_add(res, biases)

2
tensorflow/contrib/seq2seq/python/ops/sampling_decoder.py
View File

@ -113,6 +113,7 @@ class BasicSamplingDecoder(decoder.Decoder):
dtypes.int32)
def initialize(self, name=None):
with ops.name_scope("basic_sampling_decoder_initialize"):
return self._sampler.initialize() + (self._initial_state,)
def step(self, time, inputs, state):
@ -126,6 +127,7 @@ class BasicSamplingDecoder(decoder.Decoder):
Returns:
`(outputs, next_state, next_inputs, finished)`.
"""
with ops.name_scope("basic_sampling_decoder_step"):
cell_outputs, next_state = self._cell(inputs, state)
(sample_id, finished, next_inputs) = self._sampler.sample(
time=time, outputs=cell_outputs, state=next_state)

13
tensorflow/core/platform/default/build_config.bzl
View File

@ -7,7 +7,6 @@ load("//tensorflow:tensorflow.bzl", "if_not_mobile")
# configure may change the following lines
WITH_GCP_SUPPORT = False
WITH_HDFS_SUPPORT = False
WITH_XLA_SUPPORT = False
WITH_JEMALLOC = True
# Appends a suffix to a list of deps.
@ -242,15 +241,3 @@ def tf_additional_cloud_kernel_deps():
#if WITH_GCP_SUPPORT:
# deps = if_not_mobile(["//tensorflow/core:cloud_ops_op_lib"])
return deps
def tf_additional_plugin_deps():
deps = []
if WITH_XLA_SUPPORT:
deps.append("//tensorflow/compiler/jit")
return deps
def tf_additional_license_deps():
licenses = []
if WITH_XLA_SUPPORT:
licenses.append("@llvm//:LICENSE.TXT")
return licenses

17
tensorflow/core/platform/default/build_config_root.bzl
View File

@ -2,8 +2,25 @@
# The functions in this file might be referred by tensorflow.bzl. They have to
# be separate to avoid cyclic references.
WITH_XLA_SUPPORT = False
def tf_cuda_tests_tags():
return ["local"]
def tf_sycl_tests_tags():
return ["local"]
def tf_additional_plugin_deps():
deps = []
if WITH_XLA_SUPPORT:
deps.append("//tensorflow/compiler/jit")
return deps
def tf_additional_xla_deps_py():
return []
def tf_additional_license_deps():
licenses = []
if WITH_XLA_SUPPORT:
licenses.append("@llvm//:LICENSE.TXT")
return licenses

2
tensorflow/python/BUILD
View File

@ -23,7 +23,7 @@ load("//tensorflow:tensorflow.bzl", "cuda_py_tests")
load("//tensorflow/core:platform/default/build_config.bzl", "tf_proto_library")
load("//tensorflow/core:platform/default/build_config.bzl", "tf_proto_library_py")
load("//tensorflow/core:platform/default/build_config.bzl", "tf_additional_lib_deps")
load("//tensorflow/core:platform/default/build_config.bzl", "tf_additional_plugin_deps")
load("//tensorflow/core:platform/default/build_config_root.bzl", "tf_additional_plugin_deps")
load("//tensorflow/python:build_defs.bzl", "tf_gen_op_wrapper_private_py")
py_library(

29
tensorflow/tensorflow.bzl
View File

@ -12,6 +12,7 @@ load(
"//tensorflow/core:platform/default/build_config_root.bzl",
"tf_cuda_tests_tags",
"tf_sycl_tests_tags",
"tf_additional_xla_deps_py",
)
load(
"@local_config_cuda//cuda:build_defs.bzl",
@ -789,7 +790,10 @@ def py_test(deps=[], **kwargs):
**kwargs)
def tf_py_test(name, srcs, size="medium", data=[], main=None, args=[],
tags=[], shard_count=1, additional_deps=[], flaky=0):
tags=[], shard_count=1, additional_deps=[], flaky=0,
xla_enabled=False):
if xla_enabled:
additional_deps += tf_additional_xla_deps_py()
native.py_test(
name=name,
size=size,
@ -811,7 +815,8 @@ def tf_py_test(name, srcs, size="medium", data=[], main=None, args=[],
srcs_version="PY2AND3")
def cuda_py_test(name, srcs, size="medium", data=[], main=None, args=[],
shard_count=1, additional_deps=[], tags=[], flaky=0):
shard_count=1, additional_deps=[], tags=[], flaky=0,
xla_enabled=False):
test_tags = tags + tf_cuda_tests_tags()
tf_py_test(name=name,
size=size,
@ -822,10 +827,12 @@ def cuda_py_test(name, srcs, size="medium", data=[], main=None, args=[],
tags=test_tags,
shard_count=shard_count,
additional_deps=additional_deps,
flaky=flaky)
flaky=flaky,
xla_enabled=xla_enabled)
def sycl_py_test(name, srcs, size="medium", data=[], main=None, args=[],
shard_count=1, additional_deps=[], tags=[], flaky=0):
shard_count=1, additional_deps=[], tags=[], flaky=0,
xla_enabled=False):
test_tags = tags + tf_sycl_tests_tags()
tf_py_test(name=name,
size=size,
@ -836,7 +843,8 @@ def sycl_py_test(name, srcs, size="medium", data=[], main=None, args=[],
tags=test_tags,
shard_count=shard_count,
additional_deps=additional_deps,
flaky=flaky)
flaky=flaky,
xla_enabled=xla_enabled)
def py_tests(name,
srcs,
@ -845,7 +853,8 @@ def py_tests(name,
data=[],
tags=[],
shard_count=1,
prefix=""):
prefix="",
xla_enabled=False):
for src in srcs:
test_name = src.split("/")[-1].split(".")[0]
if prefix:
@ -857,13 +866,15 @@ def py_tests(name,
tags=tags,
shard_count=shard_count,
data=data,
additional_deps=additional_deps)
additional_deps=additional_deps,
xla_enabled=xla_enabled)
def cuda_py_tests(name, srcs, size="medium", additional_deps=[], data=[],
shard_count=1, tags=[], prefix=""):
shard_count=1, tags=[], prefix="", xla_enabled=False):
test_tags = tags + tf_cuda_tests_tags()
py_tests(name=name, size=size, srcs=srcs, additional_deps=additional_deps,
data=data, tags=test_tags, shard_count=shard_count,prefix=prefix)
data=data, tags=test_tags, shard_count=shard_count,prefix=prefix,
xla_enabled=xla_enabled)
# Creates a genrule named <name> for running tools/proto_text's generator to
# make the proto_text functions, for the protos passed in <srcs>.

2
tensorflow/tools/pip_package/BUILD
View File

@ -4,7 +4,7 @@
package(default_visibility = ["//visibility:private"])
load("//tensorflow:tensorflow.bzl", "transitive_hdrs")
load("//tensorflow/core:platform/default/build_config.bzl", "tf_additional_license_deps")
load("//tensorflow/core:platform/default/build_config_root.bzl", "tf_additional_license_deps")
# This returns a list of headers of all public header libraries (e.g.,
# framework, lib), and all of the transitive dependencies of those