wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/kernel_tests/data_structures/BUILD
+373
View File
@@ -0,0 +1,373 @@
# Tests of TensorFlow kernels written using the Python API.
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
tf_py_strict_test(
name = "barrier_ops_test",
size = "medium", # NOTE(ebrevdo): This test is NOT small.
srcs = ["barrier_ops_test.py"],
shard_count = 20,
tags = [
"no_mac", # TODO(b/129706424): Re-enable this test on Mac.
],
deps = [
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_py_strict_test(
name = "conditional_accumulator_test",
size = "small",
srcs = ["conditional_accumulator_test.py"],
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "dynamic_partition_op_test",
size = "medium",
srcs = ["dynamic_partition_op_test.py"],
tags = [
"multi_and_single_gpu",
],
deps = [
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_grad",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "dynamic_stitch_op_test",
size = "small",
srcs = ["dynamic_stitch_op_test.py"],
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_grad",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "dynamic_stitch_op_spmd_test",
size = "small",
srcs = ["dynamic_stitch_op_test.py"],
env = {"TF_FLAG_TF_SHAPE_DEFAULT_INT64": "true"},
main = "dynamic_stitch_op_test.py",
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_grad",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_py_strict_test(
name = "fifo_queue_test",
size = "small",
srcs = ["fifo_queue_test.py"],
tags = ["no_rocm"],
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/module",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:resource_variable_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:compat",
"//third_party/py/numpy",
],
)
tf_py_strict_test(
name = "listdiff_op_test",
size = "small",
srcs = ["listdiff_op_test.py"],
deps = [
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:compat",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "list_ops_test",
size = "small",
srcs = ["list_ops_test.py"],
grpc_enabled = True,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/ops:list_ops_gen",
"//tensorflow/python/ops:map_fn",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_py_strict_test(
name = "lookup_ops_test",
size = "small",
srcs = ["lookup_ops_test.py"],
grpc_enabled = True,
deps = [
"//tensorflow/python:tf2",
"//tensorflow/python/checkpoint:graph_view",
"//tensorflow/python/checkpoint:util",
"//tensorflow/python/client:session",
"//tensorflow/python/data/experimental/ops:counter",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:wrap_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/framework:test_ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:lookup_ops",
"//tensorflow/python/ops:lookup_ops_gen",
"//tensorflow/python/ops:map_fn",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/saved_model:load",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/trackable:asset",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/training",
"//tensorflow/python/training:saver",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:compat",
"@absl_py//absl/testing:parameterized",
],
)
# TODO(kattian): add GPU capability and change to cuda_py_test
tf_py_strict_test(
name = "map_ops_test",
size = "small",
srcs = ["map_ops_test.py"],
grpc_enabled = True,
tags = [
"no_windows", # TODO(b/192259628)
],
deps = [
"//tensorflow/python/eager:backprop",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:map_ops",
"//tensorflow/python/ops:sort_ops",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "map_stage_op_test",
size = "medium",
srcs = ["map_stage_op_test.py"],
tags = ["no_oss"], # b/124474135
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "padding_fifo_queue_test",
size = "small",
srcs = ["padding_fifo_queue_test.py"],
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_py_strict_test(
name = "priority_queue_test",
size = "medium",
srcs = ["priority_queue_test.py"],
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:nn_grad",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "stack_ops_test",
size = "small",
srcs = ["stack_ops_test.py"],
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:data_flow_ops_gen",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "stage_op_test",
size = "medium",
srcs = ["stage_op_test.py"],
deps = [
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "tensor_array_ops_test",
size = "medium",
srcs = ["tensor_array_ops_test.py"],
flaky = 1, # create_local_cluster sometimes times out.
shard_count = 10,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python:distributed_framework_test_lib",
"//tensorflow/python/client:session",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_util",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:data_flow_ops_gen",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_grad",
"//tensorflow/python/ops:tensor_array_grad",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tensorflow/python/kernel_tests/data_structures/barrier_ops_test.py
+748
View File
@@ -0,0 +1,748 @@
# 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 barrier ops."""
import time
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.framework import test_util
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.platform import test
class BarrierTest(test.TestCase):
def testConstructorWithShapes(self):
with ops.Graph().as_default():
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32),
shapes=((1, 2, 3), (8,)),
shared_name="B",
name="B")
self.assertTrue(isinstance(b.barrier_ref, tensor.Tensor))
self.assertProtoEquals("""
name:'B' op:'Barrier'
attr {
key: "capacity"
value {
i: -1
}
}
attr { key: 'component_types'
value { list { type: DT_FLOAT type: DT_FLOAT } } }
attr {
key: 'shapes'
value {
list {
shape {
dim { size: 1 } dim { size: 2 } dim { size: 3 }
}
shape {
dim { size: 8 }
}
}
}
}
attr { key: 'container' value { s: "" } }
attr { key: 'shared_name' value: { s: 'B' } }
""", b.barrier_ref.op.node_def)
@test_util.run_deprecated_v1
def testInsertMany(self):
with self.cached_session():
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shapes=((), ()), name="B")
size_t = b.ready_size()
self.assertEqual([], size_t.get_shape())
keys = [b"a", b"b", b"c"]
insert_0_op = b.insert_many(0, keys, [10.0, 20.0, 30.0])
insert_1_op = b.insert_many(1, keys, [100.0, 200.0, 300.0])
self.assertEqual(self.evaluate(size_t), [0])
insert_0_op.run()
self.assertEqual(self.evaluate(size_t), [0])
insert_1_op.run()
self.assertEqual(self.evaluate(size_t), [3])
def testInsertManyEmptyTensor(self):
with self.cached_session():
error_message = ("Empty tensors are not supported, but received shape "
r"\'\(0,\)\' at index 1")
with self.assertRaisesRegex(ValueError, error_message):
data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shapes=((1,), (0,)), name="B")
@test_util.run_deprecated_v1
def testInsertManyEmptyTensorUnknown(self):
with self.cached_session():
b = data_flow_ops.Barrier((dtypes.float32, dtypes.float32), name="B")
size_t = b.ready_size()
self.assertEqual([], size_t.get_shape())
keys = [b"a", b"b", b"c"]
insert_0_op = b.insert_many(0, keys, np.array([[], [], []], np.float32))
self.assertEqual(self.evaluate(size_t), [0])
with self.assertRaisesOpError(
".*Tensors with no elements are not supported.*"):
insert_0_op.run()
@test_util.run_deprecated_v1
def testTakeMany(self):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shapes=((), ()), name="B")
size_t = b.ready_size()
keys = [b"a", b"b", b"c"]
values_0 = [10.0, 20.0, 30.0]
values_1 = [100.0, 200.0, 300.0]
insert_0_op = b.insert_many(0, keys, values_0)
insert_1_op = b.insert_many(1, keys, values_1)
take_t = b.take_many(3)
insert_0_op.run()
insert_1_op.run()
self.assertEqual(self.evaluate(size_t), [3])
indices_val, keys_val, values_0_val, values_1_val = sess.run(
[take_t[0], take_t[1], take_t[2][0], take_t[2][1]])
self.assertAllEqual(indices_val, [-2**63] * 3)
for k, v0, v1 in zip(keys, values_0, values_1):
idx = keys_val.tolist().index(k)
self.assertEqual(values_0_val[idx], v0)
self.assertEqual(values_1_val[idx], v1)
@test_util.run_deprecated_v1
def testTakeManySmallBatch(self):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shapes=((), ()), name="B")
size_t = b.ready_size()
size_i = b.incomplete_size()
keys = [b"a", b"b", b"c", b"d"]
values_0 = [10.0, 20.0, 30.0, 40.0]
values_1 = [100.0, 200.0, 300.0, 400.0]
insert_0_op = b.insert_many(0, keys, values_0)
# Split adding of the second component into two independent operations.
# After insert_1_1_op, we'll have two ready elements in the barrier,
# 2 will still be incomplete.
insert_1_1_op = b.insert_many(1, keys[0:2], values_1[0:2]) # add "a", "b"
insert_1_2_op = b.insert_many(1, keys[2:3], values_1[2:3]) # add "c"
insert_1_3_op = b.insert_many(1, keys[3:], values_1[3:]) # add "d"
insert_empty_op = b.insert_many(0, [], [])
close_op = b.close()
close_op_final = b.close(cancel_pending_enqueues=True)
index_t, key_t, value_list_t = b.take_many(3, allow_small_batch=True)
insert_0_op.run()
insert_1_1_op.run()
close_op.run()
# Now we have a closed barrier with 2 ready elements. Running take_t
# should return a reduced batch with 2 elements only.
self.assertEqual(self.evaluate(size_i),
[2]) # assert that incomplete size = 2
self.assertEqual(self.evaluate(size_t), [2]) # assert that ready size = 2
_, keys_val, values_0_val, values_1_val = sess.run(
[index_t, key_t, value_list_t[0], value_list_t[1]])
# Check that correct values have been returned.
for k, v0, v1 in zip(keys[0:2], values_0[0:2], values_1[0:2]):
idx = keys_val.tolist().index(k)
self.assertEqual(values_0_val[idx], v0)
self.assertEqual(values_1_val[idx], v1)
# The next insert completes the element with key "c". The next take_t
# should return a batch with just 1 element.
insert_1_2_op.run()
self.assertEqual(self.evaluate(size_i),
[1]) # assert that incomplete size = 1
self.assertEqual(self.evaluate(size_t), [1]) # assert that ready size = 1
_, keys_val, values_0_val, values_1_val = sess.run(
[index_t, key_t, value_list_t[0], value_list_t[1]])
# Check that correct values have been returned.
for k, v0, v1 in zip(keys[2:3], values_0[2:3], values_1[2:3]):
idx = keys_val.tolist().index(k)
self.assertEqual(values_0_val[idx], v0)
self.assertEqual(values_1_val[idx], v1)
# Adding nothing ought to work, even if the barrier is closed.
insert_empty_op.run()
# currently keys "a" and "b" are not in the barrier, adding them
# again after it has been closed, ought to cause failure.
with self.assertRaisesOpError("is closed"):
insert_1_1_op.run()
close_op_final.run()
# These ops should fail because the barrier has now been closed with
# cancel_pending_enqueues = True.
with self.assertRaisesOpError("is closed"):
insert_empty_op.run()
with self.assertRaisesOpError("is closed"):
insert_1_3_op.run()
@test_util.run_deprecated_v1
def testUseBarrierWithShape(self):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shapes=((2, 2), (8,)), name="B")
size_t = b.ready_size()
keys = [b"a", b"b", b"c"]
values_0 = np.array(
[[[10.0] * 2] * 2, [[20.0] * 2] * 2, [[30.0] * 2] * 2], np.float32)
values_1 = np.array([[100.0] * 8, [200.0] * 8, [300.0] * 8], np.float32)
insert_0_op = b.insert_many(0, keys, values_0)
insert_1_op = b.insert_many(1, keys, values_1)
take_t = b.take_many(3)
insert_0_op.run()
insert_1_op.run()
self.assertEqual(self.evaluate(size_t), [3])
indices_val, keys_val, values_0_val, values_1_val = sess.run(
[take_t[0], take_t[1], take_t[2][0], take_t[2][1]])
self.assertAllEqual(indices_val, [-2**63] * 3)
self.assertShapeEqual(keys_val, take_t[1])
self.assertShapeEqual(values_0_val, take_t[2][0])
self.assertShapeEqual(values_1_val, take_t[2][1])
for k, v0, v1 in zip(keys, values_0, values_1):
idx = keys_val.tolist().index(k)
self.assertAllEqual(values_0_val[idx], v0)
self.assertAllEqual(values_1_val[idx], v1)
@test_util.run_deprecated_v1
def testParallelInsertMany(self):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(dtypes.float32, shapes=())
size_t = b.ready_size()
keys = [str(x).encode("ascii") for x in range(10)]
values = [float(x) for x in range(10)]
insert_ops = [b.insert_many(0, [k], [v]) for k, v in zip(keys, values)]
take_t = b.take_many(10)
self.evaluate(insert_ops)
self.assertEqual(self.evaluate(size_t), [10])
indices_val, keys_val, values_val = sess.run(
[take_t[0], take_t[1], take_t[2][0]])
self.assertAllEqual(indices_val, [-2**63 + x for x in range(10)])
for k, v in zip(keys, values):
idx = keys_val.tolist().index(k)
self.assertEqual(values_val[idx], v)
@test_util.run_deprecated_v1
def testParallelTakeMany(self):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(dtypes.float32, shapes=())
size_t = b.ready_size()
keys = [str(x).encode("ascii") for x in range(10)]
values = [float(x) for x in range(10)]
insert_op = b.insert_many(0, keys, values)
take_t = [b.take_many(1) for _ in keys]
insert_op.run()
self.assertEqual(self.evaluate(size_t), [10])
index_fetches = []
key_fetches = []
value_fetches = []
for ix_t, k_t, v_t in take_t:
index_fetches.append(ix_t)
key_fetches.append(k_t)
value_fetches.append(v_t[0])
vals = sess.run(index_fetches + key_fetches + value_fetches)
index_vals = vals[:len(keys)]
key_vals = vals[len(keys):2 * len(keys)]
value_vals = vals[2 * len(keys):]
taken_elems = []
for k, v in zip(key_vals, value_vals):
taken_elems.append((k[0], v[0]))
self.assertAllEqual(np.hstack(index_vals), [-2**63] * 10)
self.assertItemsEqual(
zip(keys, values), [(k[0], v[0]) for k, v in zip(key_vals, value_vals)])
@test_util.run_deprecated_v1
def testBlockingTakeMany(self):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(dtypes.float32, shapes=())
keys = [str(x).encode("ascii") for x in range(10)]
values = [float(x) for x in range(10)]
insert_ops = [b.insert_many(0, [k], [v]) for k, v in zip(keys, values)]
take_t = b.take_many(10)
def take():
indices_val, keys_val, values_val = sess.run(
[take_t[0], take_t[1], take_t[2][0]])
self.assertAllEqual(indices_val,
[int(x.decode("ascii")) - 2**63 for x in keys_val])
self.assertItemsEqual(zip(keys, values), zip(keys_val, values_val))
t = self.checkedThread(target=take)
t.start()
time.sleep(0.1)
for insert_op in insert_ops:
insert_op.run()
t.join()
@test_util.run_deprecated_v1
def testParallelInsertManyTakeMany(self):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.int64), shapes=((), (2,)))
num_iterations = 100
keys = [str(x) for x in range(10)]
values_0 = np.asarray(range(10), dtype=np.float32)
values_1 = np.asarray([[x + 1, x + 2] for x in range(10)], dtype=np.int64)
keys_i = lambda i: [("%d:%s" % (i, k)).encode("ascii") for k in keys]
insert_0_ops = [
b.insert_many(0, keys_i(i), values_0 + i)
for i in range(num_iterations)
]
insert_1_ops = [
b.insert_many(1, keys_i(i), values_1 + i)
for i in range(num_iterations)
]
take_ops = [b.take_many(10) for _ in range(num_iterations)]
def take(sess, i, taken):
indices_val, keys_val, values_0_val, values_1_val = sess.run([
take_ops[i][0], take_ops[i][1], take_ops[i][2][0], take_ops[i][2][1]
])
taken.append({
"indices": indices_val,
"keys": keys_val,
"values_0": values_0_val,
"values_1": values_1_val
})
def insert(sess, i):
sess.run([insert_0_ops[i], insert_1_ops[i]])
taken = []
take_threads = [
self.checkedThread(
target=take, args=(sess, i, taken)) for i in range(num_iterations)
]
insert_threads = [
self.checkedThread(
target=insert, args=(sess, i)) for i in range(num_iterations)
]
for t in take_threads:
t.start()
time.sleep(0.1)
for t in insert_threads:
t.start()
for t in take_threads:
t.join()
for t in insert_threads:
t.join()
self.assertEqual(len(taken), num_iterations)
flatten = lambda l: [item for sublist in l for item in sublist]
all_indices = sorted(flatten([t_i["indices"] for t_i in taken]))
all_keys = sorted(flatten([t_i["keys"] for t_i in taken]))
expected_keys = sorted(
flatten([keys_i(i) for i in range(num_iterations)]))
expected_indices = sorted(
flatten([-2**63 + j] * 10 for j in range(num_iterations)))
self.assertAllEqual(all_indices, expected_indices)
self.assertAllEqual(all_keys, expected_keys)
for taken_i in taken:
outer_indices_from_keys = np.array(
[int(k.decode("ascii").split(":")[0]) for k in taken_i["keys"]])
inner_indices_from_keys = np.array(
[int(k.decode("ascii").split(":")[1]) for k in taken_i["keys"]])
self.assertAllEqual(taken_i["values_0"],
outer_indices_from_keys + inner_indices_from_keys)
expected_values_1 = np.vstack(
(1 + outer_indices_from_keys + inner_indices_from_keys,
2 + outer_indices_from_keys + inner_indices_from_keys)).T
self.assertAllEqual(taken_i["values_1"], expected_values_1)
@test_util.run_deprecated_v1
def testClose(self):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shapes=((), ()), name="B")
size_t = b.ready_size()
incomplete_t = b.incomplete_size()
keys = [b"a", b"b", b"c"]
values_0 = [10.0, 20.0, 30.0]
values_1 = [100.0, 200.0, 300.0]
insert_0_op = b.insert_many(0, keys, values_0)
insert_1_op = b.insert_many(1, keys, values_1)
close_op = b.close()
fail_insert_op = b.insert_many(0, ["f"], [60.0])
take_t = b.take_many(3)
take_too_many_t = b.take_many(4)
self.assertEqual(self.evaluate(size_t), [0])
self.assertEqual(self.evaluate(incomplete_t), [0])
insert_0_op.run()
self.assertEqual(self.evaluate(size_t), [0])
self.assertEqual(self.evaluate(incomplete_t), [3])
close_op.run()
# This op should fail because the barrier is closed.
with self.assertRaisesOpError("is closed"):
fail_insert_op.run()
# This op should succeed because the barrier has not canceled
# pending enqueues
insert_1_op.run()
self.assertEqual(self.evaluate(size_t), [3])
self.assertEqual(self.evaluate(incomplete_t), [0])
# This op should fail because the barrier is closed.
with self.assertRaisesOpError("is closed"):
fail_insert_op.run()
# This op should fail because we requested more elements than are
# available in incomplete + ready queue.
with self.assertRaisesOpError(r"is closed and has insufficient elements "
r"\(requested 4, total size 3\)"):
sess.run(take_too_many_t[0]) # Sufficient to request just the indices
# This op should succeed because there are still completed elements
# to process.
indices_val, keys_val, values_0_val, values_1_val = sess.run(
[take_t[0], take_t[1], take_t[2][0], take_t[2][1]])
self.assertAllEqual(indices_val, [-2**63] * 3)
for k, v0, v1 in zip(keys, values_0, values_1):
idx = keys_val.tolist().index(k)
self.assertEqual(values_0_val[idx], v0)
self.assertEqual(values_1_val[idx], v1)
# This op should fail because there are no more completed elements and
# the queue is closed.
with self.assertRaisesOpError("is closed and has insufficient elements"):
sess.run(take_t[0])
@test_util.run_deprecated_v1
def testCancel(self):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shapes=((), ()), name="B")
size_t = b.ready_size()
incomplete_t = b.incomplete_size()
keys = [b"a", b"b", b"c"]
values_0 = [10.0, 20.0, 30.0]
values_1 = [100.0, 200.0, 300.0]
insert_0_op = b.insert_many(0, keys, values_0)
insert_1_op = b.insert_many(1, keys[0:2], values_1[0:2])
insert_2_op = b.insert_many(1, keys[2:], values_1[2:])
cancel_op = b.close(cancel_pending_enqueues=True)
fail_insert_op = b.insert_many(0, ["f"], [60.0])
take_t = b.take_many(2)
take_too_many_t = b.take_many(3)
self.assertEqual(self.evaluate(size_t), [0])
insert_0_op.run()
insert_1_op.run()
self.assertEqual(self.evaluate(size_t), [2])
self.assertEqual(self.evaluate(incomplete_t), [1])
cancel_op.run()
# This op should fail because the queue is closed.
with self.assertRaisesOpError("is closed"):
fail_insert_op.run()
# This op should fail because the queue is canceled.
with self.assertRaisesOpError("is closed"):
insert_2_op.run()
# This op should fail because we requested more elements than are
# available in incomplete + ready queue.
with self.assertRaisesOpError(r"is closed and has insufficient elements "
r"\(requested 3, total size 2\)"):
sess.run(take_too_many_t[0]) # Sufficient to request just the indices
# This op should succeed because there are still completed elements
# to process.
indices_val, keys_val, values_0_val, values_1_val = sess.run(
[take_t[0], take_t[1], take_t[2][0], take_t[2][1]])
self.assertAllEqual(indices_val, [-2**63] * 2)
for k, v0, v1 in zip(keys[0:2], values_0[0:2], values_1[0:2]):
idx = keys_val.tolist().index(k)
self.assertEqual(values_0_val[idx], v0)
self.assertEqual(values_1_val[idx], v1)
# This op should fail because there are no more completed elements and
# the queue is closed.
with self.assertRaisesOpError("is closed and has insufficient elements"):
sess.run(take_t[0])
def _testClosedEmptyBarrierTakeManyAllowSmallBatchRaises(self, cancel):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shapes=((), ()), name="B")
take_t = b.take_many(1, allow_small_batch=True)
self.evaluate(b.close(cancel))
with self.assertRaisesOpError("is closed and has insufficient elements"):
self.evaluate(take_t)
@test_util.run_deprecated_v1
def testClosedEmptyBarrierTakeManyAllowSmallBatchRaises(self):
self._testClosedEmptyBarrierTakeManyAllowSmallBatchRaises(cancel=False)
self._testClosedEmptyBarrierTakeManyAllowSmallBatchRaises(cancel=True)
def _testParallelInsertManyTakeManyCloseHalfwayThrough(self, cancel):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.int64), shapes=((), (2,)))
num_iterations = 50
keys = [str(x) for x in range(10)]
values_0 = np.asarray(range(10), dtype=np.float32)
values_1 = np.asarray([[x + 1, x + 2] for x in range(10)], dtype=np.int64)
keys_i = lambda i: [("%d:%s" % (i, k)).encode("ascii") for k in keys]
insert_0_ops = [
b.insert_many(0, keys_i(i), values_0 + i)
for i in range(num_iterations)
]
insert_1_ops = [
b.insert_many(1, keys_i(i), values_1 + i)
for i in range(num_iterations)
]
take_ops = [b.take_many(10) for _ in range(num_iterations)]
close_op = b.close(cancel_pending_enqueues=cancel)
def take(sess, i, taken):
try:
indices_val, unused_keys_val, unused_val_0, unused_val_1 = sess.run([
take_ops[i][0], take_ops[i][1], take_ops[i][2][0],
take_ops[i][2][1]
])
taken.append(len(indices_val))
except errors_impl.OutOfRangeError:
taken.append(0)
def insert(sess, i):
try:
sess.run([insert_0_ops[i], insert_1_ops[i]])
except errors_impl.CancelledError:
pass
taken = []
take_threads = [
self.checkedThread(
target=take, args=(sess, i, taken)) for i in range(num_iterations)
]
insert_threads = [
self.checkedThread(
target=insert, args=(sess, i)) for i in range(num_iterations)
]
first_half_insert_threads = insert_threads[:num_iterations // 2]
second_half_insert_threads = insert_threads[num_iterations // 2:]
for t in take_threads:
t.start()
for t in first_half_insert_threads:
t.start()
for t in first_half_insert_threads:
t.join()
close_op.run()
for t in second_half_insert_threads:
t.start()
for t in take_threads:
t.join()
for t in second_half_insert_threads:
t.join()
self.assertEqual(
sorted(taken),
[0] * (num_iterations // 2) + [10] * (num_iterations // 2))
@test_util.run_deprecated_v1
def testParallelInsertManyTakeManyCloseHalfwayThrough(self):
self._testParallelInsertManyTakeManyCloseHalfwayThrough(cancel=False)
@test_util.run_deprecated_v1
def testParallelInsertManyTakeManyCancelHalfwayThrough(self):
self._testParallelInsertManyTakeManyCloseHalfwayThrough(cancel=True)
def _testParallelPartialInsertManyTakeManyCloseHalfwayThrough(self, cancel):
with self.cached_session() as sess:
b = data_flow_ops.Barrier(
(dtypes.float32, dtypes.int64), shapes=((), (2,)))
num_iterations = 100
keys = [str(x) for x in range(10)]
values_0 = np.asarray(range(10), dtype=np.float32)
values_1 = np.asarray([[x + 1, x + 2] for x in range(10)], dtype=np.int64)
keys_i = lambda i: [("%d:%s" % (i, k)).encode("ascii") for k in keys]
insert_0_ops = [
b.insert_many(
0, keys_i(i), values_0 + i, name="insert_0_%d" % i)
for i in range(num_iterations)
]
close_op = b.close(cancel_pending_enqueues=cancel)
take_ops = [
b.take_many(
10, name="take_%d" % i) for i in range(num_iterations)
]
# insert_1_ops will only run after closure
insert_1_ops = [
b.insert_many(
1, keys_i(i), values_1 + i, name="insert_1_%d" % i)
for i in range(num_iterations)
]
def take(sess, i, taken):
if cancel:
try:
indices_val, unused_keys_val, unused_val_0, unused_val_1 = sess.run(
[
take_ops[i][0], take_ops[i][1], take_ops[i][2][0],
take_ops[i][2][1]
])
taken.append(len(indices_val))
except errors_impl.OutOfRangeError:
taken.append(0)
else:
indices_val, unused_keys_val, unused_val_0, unused_val_1 = sess.run([
take_ops[i][0], take_ops[i][1], take_ops[i][2][0],
take_ops[i][2][1]
])
taken.append(len(indices_val))
def insert_0(sess, i):
insert_0_ops[i].run(session=sess)
def insert_1(sess, i):
if cancel:
try:
insert_1_ops[i].run(session=sess)
except errors_impl.CancelledError:
pass
else:
insert_1_ops[i].run(session=sess)
taken = []
take_threads = [
self.checkedThread(
target=take, args=(sess, i, taken)) for i in range(num_iterations)
]
insert_0_threads = [
self.checkedThread(
target=insert_0, args=(sess, i)) for i in range(num_iterations)
]
insert_1_threads = [
self.checkedThread(
target=insert_1, args=(sess, i)) for i in range(num_iterations)
]
for t in insert_0_threads:
t.start()
for t in insert_0_threads:
t.join()
for t in take_threads:
t.start()
close_op.run()
for t in insert_1_threads:
t.start()
for t in take_threads:
t.join()
for t in insert_1_threads:
t.join()
if cancel:
self.assertEqual(taken, [0] * num_iterations)
else:
self.assertEqual(taken, [10] * num_iterations)
@test_util.run_deprecated_v1
def testParallelPartialInsertManyTakeManyCloseHalfwayThrough(self):
self._testParallelPartialInsertManyTakeManyCloseHalfwayThrough(cancel=False)
@test_util.run_deprecated_v1
def testParallelPartialInsertManyTakeManyCancelHalfwayThrough(self):
self._testParallelPartialInsertManyTakeManyCloseHalfwayThrough(cancel=True)
@test_util.run_deprecated_v1
def testIncompatibleSharedBarrierErrors(self):
with self.cached_session():
# Do component types and shapes.
b_a_1 = data_flow_ops.Barrier(
(dtypes.float32,), shapes=(()), shared_name="b_a")
b_a_2 = data_flow_ops.Barrier(
(dtypes.int32,), shapes=(()), shared_name="b_a")
self.evaluate(b_a_1.barrier_ref)
with self.assertRaisesOpError("component types"):
self.evaluate(b_a_2.barrier_ref)
b_b_1 = data_flow_ops.Barrier(
(dtypes.float32,), shapes=(()), shared_name="b_b")
b_b_2 = data_flow_ops.Barrier(
(dtypes.float32, dtypes.int32), shapes=((), ()), shared_name="b_b")
self.evaluate(b_b_1.barrier_ref)
with self.assertRaisesOpError("component types"):
self.evaluate(b_b_2.barrier_ref)
b_c_1 = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32),
shapes=((2, 2), (8,)),
shared_name="b_c")
b_c_2 = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shared_name="b_c")
self.evaluate(b_c_1.barrier_ref)
with self.assertRaisesOpError("component shapes"):
self.evaluate(b_c_2.barrier_ref)
b_d_1 = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32), shapes=((), ()), shared_name="b_d")
b_d_2 = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32),
shapes=((2, 2), (8,)),
shared_name="b_d")
self.evaluate(b_d_1.barrier_ref)
with self.assertRaisesOpError("component shapes"):
self.evaluate(b_d_2.barrier_ref)
b_e_1 = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32),
shapes=((2, 2), (8,)),
shared_name="b_e")
b_e_2 = data_flow_ops.Barrier(
(dtypes.float32, dtypes.float32),
shapes=((2, 5), (8,)),
shared_name="b_e")
self.evaluate(b_e_1.barrier_ref)
with self.assertRaisesOpError("component shapes"):
self.evaluate(b_e_2.barrier_ref)
if __name__ == "__main__":
test.main()
tensorflow/python/kernel_tests/data_structures/conditional_accumulator_test.py
+521
View File
@@ -0,0 +1,521 @@
# Copyright 2015 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.
# ==============================================================================
import time
import numpy as np
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes as dtypes_lib
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
# from functools import reduce
class ConditionalAccumulatorTest(test.TestCase):
def testConstructorWithInvalidArg(self):
with ops.Graph().as_default():
with self.assertRaises(ValueError):
data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", reduction_type="Invalid")
@test_util.run_deprecated_v1
def testAccumulatorSizeEmpty(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(dtypes_lib.float32, name="Q")
self.assertEqual(q.num_accumulated().eval(), 0)
@test_util.run_deprecated_v1
def testAccumulatorSetGlobalStep(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
set_global_step_op = q.set_global_step(1)
set_global_step_op.run()
@test_util.run_deprecated_v1
def testAccumulatorApplyGradFloat32(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
accum_op = q.apply_grad((10.0,))
accum_op.run()
@test_util.run_deprecated_v1
def testDtypes(self):
with self.cached_session() as sess:
dtypes = [dtypes_lib.float16, dtypes_lib.float32, dtypes_lib.float64]
for i in range(len(dtypes)):
dtype = dtypes[i]
q = data_flow_ops.ConditionalAccumulator(
dtype, shape=tensor_shape.TensorShape([1]))
elems = np.arange(10).astype(dtype.as_numpy_dtype)
for e in elems:
q.apply_grad((e,)).run()
result = self.evaluate(q.take_grad(1))
self.assertEqual(sum(elems) / len(elems), result)
@test_util.run_deprecated_v1
def testAccumulatorMultipleAccumulators(self):
with self.cached_session():
q_f32_0 = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
q_f32_1 = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
q_f16_0 = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float16, name="Q", shape=tensor_shape.TensorShape([1]))
q_f16_1 = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float16, name="Q", shape=tensor_shape.TensorShape([1]))
accums = [q_f16_0, q_f16_1, q_f32_0, q_f32_1]
for i in range(len(accums)):
accums[i].apply_grad((i + 10.0,)).run()
for i in range(len(accums)):
result = accums[i].take_grad(1).eval()
self.assertEqual(result, i + 10.0)
@test_util.run_deprecated_v1
def testAccumulatorApplyAndTakeGradWithShape(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=(3, 2))
elems = [[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
[[10.0, 20.0], [30.0, 40.0], [50.0, 60.0]]]
elems_ave = [[(a + b) / len(elems) for a, b in zip(x, y)]
for x, y in zip(elems[0], elems[1])]
accum_ops = [q.apply_grad(x) for x in elems]
takeg_t = q.take_grad(1)
for accum_op in accum_ops:
accum_op.run()
is_all_equal = True
val = self.evaluate(takeg_t)
for i in range(len(val)):
for j in range(len(val[i])):
is_all_equal &= (val[i][j] == elems_ave[i][j])
self.assertTrue(is_all_equal)
@test_util.run_deprecated_v1
def testAccumulatorApplyGradWithWrongShape(self):
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=(3, 2))
with self.assertRaises(ValueError):
q.apply_grad([[1.0, 2.0], [3.0, 4.0]])
with self.assertRaises(ValueError):
q.apply_grad([[1.0], [2.0], [3.0]])
@test_util.run_deprecated_v1
def testAccumulatorDynamicShape(self):
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=None)
x = array_ops.placeholder(dtypes_lib.float32)
accum_op = q.apply_grad(x)
elems = [[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
[[10.0, 20.0], [30.0, 40.0], [50.0, 60.0]]]
elems_ave = [[(a + b) / len(elems) for a, b in zip(c, d)]
for c, d in zip(elems[0], elems[1])]
takeg_t = q.take_grad(1)
for elem in elems:
sess.run(accum_op, feed_dict={x: elem})
is_all_equal = True
val = self.evaluate(takeg_t)
for i in range(len(val)):
for j in range(len(val[i])):
is_all_equal &= (val[i][j] == elems_ave[i][j])
self.assertTrue(is_all_equal)
@test_util.run_v1_only("b/120545219")
def testAccumulatorWrongDynamicShape(self):
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=None)
x = array_ops.placeholder(dtypes_lib.float32)
accum_op = q.apply_grad(x)
# First successful apply_grad determines shape
sess.run(accum_op, feed_dict={x: [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]})
with self.assertRaises(errors_impl.InvalidArgumentError):
sess.run(accum_op, feed_dict={x: [[1.0, 2.0], [3.0, 4.0]]})
with self.assertRaises(errors_impl.InvalidArgumentError):
sess.run(accum_op, feed_dict={x: [[1.0], [2.0], [3.0]]})
@test_util.run_deprecated_v1
def testAccumulatorSizeAfterApplyGrad(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
accum_op = q.apply_grad((10.0,))
self.assertEqual(q.num_accumulated().eval(), 0)
accum_op.run()
self.assertEqual(q.num_accumulated().eval(), 1)
accum_op.run()
self.assertEqual(q.num_accumulated().eval(), 2)
@test_util.run_deprecated_v1
def testAccumulatorSizeAfterApplyGradAndTakeGrad(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
accum_op = q.apply_grad((10.0,))
extract_t = q.take_grad(2)
# Applying gradient multiple times to increase size from 0 to 2.
self.assertEqual(q.num_accumulated().eval(), 0)
accum_op.run()
self.assertEqual(q.num_accumulated().eval(), 1)
accum_op.run()
self.assertEqual(q.num_accumulated().eval(), 2)
# Extract will reduce size to 0
extract_t.op.run()
self.assertEqual(q.num_accumulated().eval(), 0)
# Take gradients always sets the size back to 0 if successful.
accum_op = q.apply_grad((10.0,), local_step=1)
accum_op.run()
accum_op.run()
accum_op.run()
accum_op.run()
self.assertEqual(q.num_accumulated().eval(), 4)
extract_t.op.run()
self.assertEqual(q.num_accumulated().eval(), 0)
@test_util.run_deprecated_v1
def testAccumulatorTakeGradMean(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
elems = [10.0, 20.0]
accum_ops = [q.apply_grad((x,), local_step=0) for x in elems]
takeg_t = q.take_grad(1)
for accum_op in accum_ops:
accum_op.run()
val = self.evaluate(takeg_t)
self.assertEqual(15.0, val)
accum_ops = [q.apply_grad((x,), local_step=1) for x in elems]
takeg_t = q.take_grad(constant_op.constant(1))
for accum_op in accum_ops:
accum_op.run()
val = self.evaluate(takeg_t)
self.assertEqual(15.0, val)
@test_util.run_deprecated_v1
def testAccumulatorTakeGradSum(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32,
name="Q",
shape=tensor_shape.TensorShape([1]),
reduction_type="SUM")
elems = [10.0, 20.0]
accum_ops = [q.apply_grad((x,), local_step=0) for x in elems]
takeg_t = q.take_grad(1)
for accum_op in accum_ops:
accum_op.run()
val = self.evaluate(takeg_t)
self.assertEqual(30.0, val)
accum_ops = [q.apply_grad((x,), local_step=1) for x in elems]
takeg_t = q.take_grad(constant_op.constant(1))
for accum_op in accum_ops:
accum_op.run()
val = self.evaluate(takeg_t)
self.assertEqual(30.0, val)
@test_util.run_deprecated_v1
def testAccumulatorTakeGradInvalidReductionType(self):
with self.assertRaises(ValueError):
data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32,
name="Q",
shape=tensor_shape.TensorShape([1]),
reduction_type="Invalid")
@test_util.run_v1_only("b/120545219")
def testAccumulatorInvalidTakeGrad(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
elems = [10.0, 20.0]
accum_ops = [q.apply_grad((x,)) for x in elems]
takeg_t = q.take_grad(-1)
for accum_op in accum_ops:
accum_op.run()
with self.assertRaises(errors_impl.InvalidArgumentError):
self.evaluate(takeg_t)
@test_util.run_deprecated_v1
def testAccumulatorRepeatedTakeGradMean(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
elems = [10.0, 20.0]
elems_ave = sum(elems) / len(elems)
accum_ops = [q.apply_grad((x,), local_step=0) for x in elems]
takeg_t = q.take_grad(1)
for accum_op in accum_ops:
accum_op.run()
val = self.evaluate(takeg_t)
self.assertEqual(elems_ave, val)
elems = [20.0, 30.0]
elems_ave = sum(elems) / len(elems)
accum_ops = [q.apply_grad((x,), local_step=1) for x in elems]
takeg_t = q.take_grad(1)
for accum_op in accum_ops:
accum_op.run()
val = self.evaluate(takeg_t)
self.assertEqual(elems_ave + 0.0, val)
@test_util.run_deprecated_v1
def testAccumulatorRepeatedTakeGradSum(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32,
name="Q",
shape=tensor_shape.TensorShape([1]),
reduction_type="SUM")
elems = [10.0, 20.0]
elems_sum = 30.0
accum_ops = [q.apply_grad((x,), local_step=0) for x in elems]
takeg_t = q.take_grad(1)
for accum_op in accum_ops:
accum_op.run()
val = self.evaluate(takeg_t)
self.assertEqual(elems_sum, val)
elems = [20.0, 30.0]
elems_sum = 50.0
accum_ops = [q.apply_grad((x,), local_step=1) for x in elems]
takeg_t = q.take_grad(1)
for accum_op in accum_ops:
accum_op.run()
val = self.evaluate(takeg_t)
self.assertEqual(elems_sum, val)
@test_util.run_deprecated_v1
def testAccumulatorIncrementGlobalStep(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
global_step = variables.Variable(0, name="global_step")
new_global_step = math_ops.add(global_step, 1)
inc_global_step = state_ops.assign(global_step, new_global_step)
set_global_step_op = q.set_global_step(new_global_step)
self.evaluate(variables.global_variables_initializer())
for _ in range(3):
set_global_step_op.run()
self.evaluate(inc_global_step)
@test_util.run_deprecated_v1
def testAccumulatorSetGlobalStepPreventsAccumulation(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
local_steps = range(1000, 1005)
accum_ops = [q.apply_grad((0.0 + x,), local_step=x) for x in local_steps]
for ls in local_steps:
set_global_step_op = q.set_global_step(ls)
set_global_step_op.run()
for accum_op in accum_ops:
accum_op.run()
takeg_t = q.take_grad(1)
val = self.evaluate(takeg_t)
self.assertEqual(0.0 + sum(x for x in local_steps
if x >= ls) / sum(1 for x in local_steps
if x >= ls), val)
@test_util.run_v1_only("b/120545219")
def testParallelApplyGrad(self):
# We need each thread to keep its own device stack or the device scopes
# won't be properly nested.
ops.get_default_graph().switch_to_thread_local()
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
elems = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0]
accum_ops = [q.apply_grad((x,), local_step=0) for x in elems]
takeg_t = q.take_grad(1)
def apply_grad(accum_op):
self.evaluate(accum_op)
threads = [
self.checkedThread(
target=apply_grad, args=(o,)) for o in accum_ops
]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
val = self.evaluate(takeg_t)
self.assertEqual(val, sum(elems) / len(elems))
@test_util.run_v1_only("b/120545219")
def testParallelTakeGrad(self):
# We need each thread to keep its own device stack or the device scopes
# won't be properly nested.
ops.get_default_graph().switch_to_thread_local()
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
elems = [e for e in range(10)]
accum_ops = [q.apply_grad((np.float32(e),), local_step=e) for e in elems]
takeg_t = q.take_grad(1)
def apply_grad():
for accum_op in accum_ops:
time.sleep(1.0)
self.evaluate(accum_op)
apply_grad_thread = self.checkedThread(target=apply_grad)
results = []
def take_grad():
results.append(self.evaluate(takeg_t))
threads = [self.checkedThread(target=take_grad) for _ in range(10)]
for thread in threads:
thread.start()
apply_grad_thread.start()
for thread in threads:
thread.join()
apply_grad_thread.join()
self.assertItemsEqual(elems, results)
@test_util.run_v1_only("b/120545219")
def testAccumulatorApplyAndBlockingTake(self):
# We need each thread to keep its own device stack or the device scopes
# won't be properly nested.
ops.get_default_graph().switch_to_thread_local()
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
elems = [10.0, 20.0, 30.0]
elems_ave = sum(elems) / len(elems)
accum_ops = [q.apply_grad((x,), local_step=0) for x in elems]
takeg_t = q.take_grad(3)
def apply_grad():
time.sleep(1.0)
for accum_op in accum_ops:
self.evaluate(accum_op)
return_array = []
def take_grad():
return_array.append(self.evaluate(takeg_t))
accum_thread = self.checkedThread(target=apply_grad)
takeg_thread = self.checkedThread(target=take_grad)
accum_thread.start()
takeg_thread.start()
accum_thread.join()
takeg_thread.join()
self.assertEqual([elems_ave], return_array)
def _blocking_takeg(self, sess, takeg_op):
with self.assertRaisesOpError("was cancelled"):
self.evaluate(takeg_op)
@test_util.run_v1_only("b/120545219")
def testAccumulatorCancel(self):
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
dtypes_lib.float32, name="Q", shape=tensor_shape.TensorShape([1]))
takeg_t = q.take_grad(1)
takeg_thread = self.checkedThread(
self._blocking_takeg, args=(sess, takeg_t))
takeg_thread.start()
time.sleep(1.0)
sess.close() # Will cancel blocked operation
takeg_thread.join()
if __name__ == "__main__":
test.main()
tensorflow/python/kernel_tests/data_structures/dynamic_partition_op_test.py
+373
View File
@@ -0,0 +1,373 @@
# Copyright 2015 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 the DynamicPartition op."""
import unittest
from absl.testing import parameterized
import numpy as np
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import gradients_impl
import tensorflow.python.ops.data_flow_grad # pylint: disable=unused-import
from tensorflow.python.platform import test
class DynamicPartitionTest(test.TestCase, parameterized.TestCase):
@parameterized.parameters(dtypes.float32, dtypes.bfloat16)
@test_util.run_deprecated_v1
def testSimpleOneDimensional(self, dtype):
with self.session():
data = constant_op.constant([0, 13, 2, 39, 4, 17], dtype=dtype)
indices = constant_op.constant([0, 0, 2, 3, 2, 1])
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=4)
partition_vals = self.evaluate(partitions)
self.assertEqual(4, len(partition_vals))
self.assertAllEqual([0, 13], partition_vals[0])
self.assertAllEqual([17], partition_vals[1])
self.assertAllEqual([2, 4], partition_vals[2])
self.assertAllEqual([39], partition_vals[3])
# Vector data input to DynamicPartition results in
# `num_partitions` vectors of unknown length.
self.assertEqual([None], partitions[0].get_shape().as_list())
self.assertEqual([None], partitions[1].get_shape().as_list())
self.assertEqual([None], partitions[2].get_shape().as_list())
self.assertEqual([None], partitions[3].get_shape().as_list())
@test_util.run_deprecated_v1
def testSimpleTwoDimensional(self):
with self.session():
data = constant_op.constant([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11],
[12, 13, 14], [15, 16, 17]],
dtype=dtypes.float32)
indices = constant_op.constant([0, 0, 2, 3, 2, 1])
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=4)
partition_vals = self.evaluate(partitions)
self.assertEqual(4, len(partition_vals))
self.assertAllEqual([[0, 1, 2], [3, 4, 5]], partition_vals[0])
self.assertAllEqual([[15, 16, 17]], partition_vals[1])
self.assertAllEqual([[6, 7, 8], [12, 13, 14]], partition_vals[2])
self.assertAllEqual([[9, 10, 11]], partition_vals[3])
# Vector data input to DynamicPartition results in
# `num_partitions` matrices with an unknown number of rows, and 3 columns.
self.assertEqual([None, 3], partitions[0].get_shape().as_list())
self.assertEqual([None, 3], partitions[1].get_shape().as_list())
self.assertEqual([None, 3], partitions[2].get_shape().as_list())
self.assertEqual([None, 3], partitions[3].get_shape().as_list())
def testLargeOneDimensional(self):
num = 100000
data_list = [x for x in range(num)]
indices_list = [x % 2 for x in range(num)]
part1 = [x for x in range(num) if x % 2 == 0]
part2 = [x for x in range(num) if x % 2 == 1]
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.float32)
indices = constant_op.constant(indices_list, dtype=dtypes.int32)
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=2)
partition_vals = self.evaluate(partitions)
self.assertEqual(2, len(partition_vals))
self.assertAllEqual(part1, partition_vals[0])
self.assertAllEqual(part2, partition_vals[1])
def testLargeTwoDimensional(self):
rows = 100000
cols = 100
data_list = [None] * rows
for i in range(rows):
data_list[i] = [i for _ in range(cols)]
num_partitions = 97
indices_list = [(i**2) % num_partitions for i in range(rows)]
parts = [[] for _ in range(num_partitions)]
for i in range(rows):
parts[(i**2) % num_partitions].append(data_list[i])
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.float32)
indices = constant_op.constant(indices_list, dtype=dtypes.int32)
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=num_partitions)
partition_vals = self.evaluate(partitions)
self.assertEqual(num_partitions, len(partition_vals))
for i in range(num_partitions):
# reshape because of empty parts
parts_np = np.array(parts[i], dtype=np.float64).reshape(-1, cols)
self.assertAllEqual(parts_np, partition_vals[i])
def testSimpleComplex(self):
data_list = [1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j]
indices_list = [1, 0, 1, 0]
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.complex64)
indices = constant_op.constant(indices_list, dtype=dtypes.int32)
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=2)
partition_vals = self.evaluate(partitions)
self.assertEqual(2, len(partition_vals))
self.assertAllEqual([3 + 4j, 7 + 8j], partition_vals[0])
self.assertAllEqual([1 + 2j, 5 + 6j], partition_vals[1])
def testScalarPartitions(self):
data_list = [10, 13, 12, 11]
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.float64)
indices = 3
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=4)
partition_vals = self.evaluate(partitions)
self.assertEqual(4, len(partition_vals))
self.assertAllEqual(
np.array([], dtype=np.float64).reshape(-1, 4), partition_vals[0])
self.assertAllEqual(
np.array([], dtype=np.float64).reshape(-1, 4), partition_vals[1])
self.assertAllEqual(
np.array([], dtype=np.float64).reshape(-1, 4), partition_vals[2])
self.assertAllEqual(
np.array([10, 13, 12, 11], dtype=np.float64).reshape(-1, 4),
partition_vals[3])
@test_util.run_deprecated_v1
def testHigherRank(self):
np.random.seed(7)
with self.session() as sess:
for n in 2, 3:
for shape in (4,), (4, 5), (4, 5, 2):
partitions = np.random.randint(n, size=np.prod(shape)).reshape(shape)
for extra_shape in (), (6,), (6, 7):
data = np.random.randn(*(shape + extra_shape))
partitions_t = constant_op.constant(partitions, dtype=dtypes.int32)
data_t = constant_op.constant(data)
outputs = data_flow_ops.dynamic_partition(
data_t, partitions_t, num_partitions=n)
self.assertEqual(n, len(outputs))
outputs_val = self.evaluate(outputs)
for i, output in enumerate(outputs_val):
self.assertAllEqual(output, data[partitions == i])
# Test gradients
outputs_grad = [7 * output for output in outputs_val]
grads = gradients_impl.gradients(outputs, [data_t, partitions_t],
outputs_grad)
self.assertEqual(grads[1], None) # Partitions has no gradients
self.assertAllEqual(7 * data, sess.run(grads[0]))
def testEmptyParts(self):
data_list = [1, 2, 3, 4]
indices_list = [1, 3, 1, 3]
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.float32)
indices = constant_op.constant(indices_list, dtype=dtypes.int32)
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=4)
partition_vals = self.evaluate(partitions)
self.assertEqual(4, len(partition_vals))
self.assertAllEqual([], partition_vals[0])
self.assertAllEqual([1, 3], partition_vals[1])
self.assertAllEqual([], partition_vals[2])
self.assertAllEqual([2, 4], partition_vals[3])
def testEmptyDataTwoDimensional(self):
data_list = [[], []]
indices_list = [0, 1]
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.float32)
indices = constant_op.constant(indices_list, dtype=dtypes.int32)
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=3)
partition_vals = self.evaluate(partitions)
self.assertEqual(3, len(partition_vals))
self.assertAllEqual([[]], partition_vals[0])
self.assertAllEqual([[]], partition_vals[1])
self.assertAllEqual(
np.array([], dtype=np.float64).reshape(0, 0), partition_vals[2])
def testEmptyPartitions(self):
data_list = []
indices_list = []
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.float32)
indices = constant_op.constant(indices_list, dtype=dtypes.int32)
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=2)
partition_vals = self.evaluate(partitions)
self.assertEqual(2, len(partition_vals))
self.assertAllEqual([], partition_vals[0])
self.assertAllEqual([], partition_vals[1])
@unittest.skip("Fails on windows.")
def testGPUTooManyParts(self):
# This test only makes sense on the GPU. There we do not check
# for errors. In this case, we should discard all but the first
# num_partitions indices.
if not test.is_gpu_available():
return
data_list = [1, 2, 3, 4, 5, 6]
indices_list = [6, 5, 4, 3, 1, 0]
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.float32)
indices = constant_op.constant(indices_list, dtype=dtypes.int32)
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=2)
partition_vals = self.evaluate(partitions)
self.assertEqual(2, len(partition_vals))
self.assertAllEqual([6], partition_vals[0])
self.assertAllEqual([5], partition_vals[1])
@unittest.skip("Fails on windows.")
def testGPUPartsTooLarge(self):
# This test only makes sense on the GPU. There we do not check
# for errors. In this case, we should discard all the values
# larger than num_partitions.
if not test.is_gpu_available():
return
data_list = [1, 2, 3, 4, 5, 6]
indices_list = [10, 11, 2, 12, 0, 1000]
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.float32)
indices = constant_op.constant(indices_list, dtype=dtypes.int32)
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=5)
partition_vals = self.evaluate(partitions)
self.assertEqual(5, len(partition_vals))
self.assertAllEqual([5], partition_vals[0])
self.assertAllEqual([], partition_vals[1])
self.assertAllEqual([3], partition_vals[2])
self.assertAllEqual([], partition_vals[3])
self.assertAllEqual([], partition_vals[4])
@unittest.skip("Fails on windows.")
def testGPUAllIndicesBig(self):
# This test only makes sense on the GPU. There we do not check
# for errors. In this case, we should discard all the values
# and have an empty output.
if not test.is_gpu_available():
return
data_list = [1.1, 2.1, 3.1, 4.1, 5.1, 6.1]
indices_list = [90, 70, 60, 100, 110, 40]
with self.session():
data = constant_op.constant(data_list, dtype=dtypes.float32)
indices = constant_op.constant(indices_list, dtype=dtypes.int32)
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=40)
partition_vals = self.evaluate(partitions)
self.assertEqual(40, len(partition_vals))
for i in range(40):
self.assertAllEqual([], partition_vals[i])
@test_util.run_deprecated_v1
@test_util.no_xla_auto_jit("xla doesn't raise out-of-range exceptions")
def testErrorIndexOutOfRange(self):
# GPU kernels don't throw exceptions.
with self.cached_session(use_gpu=False):
data = constant_op.constant([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11],
[12, 13, 14]])
indices = constant_op.constant([0, 2, 99, 2, 2])
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=4)
with self.assertRaisesOpError(r"partitions\[2\] = 99 is not in \[0, 4\)"):
self.evaluate(partitions)
@test_util.run_deprecated_v1
@test_util.no_xla_auto_jit("xla doesn't raise out-of-range exceptions")
def testScalarIndexOutOfRange(self):
# GPU kernels don't throw exceptions.
with self.cached_session(use_gpu=False):
bad = 17
data = np.zeros(5)
partitions = data_flow_ops.dynamic_partition(data, bad, num_partitions=7)
with self.assertRaisesOpError(r"partitions = 17 is not in \[0, 7\)"):
self.evaluate(partitions)
@test_util.run_deprecated_v1
@test_util.no_xla_auto_jit("xla doesn't raise out-of-range exceptions")
def testHigherRankIndexOutOfRange(self):
# GPU kernels don't throw exceptions.
with self.cached_session(use_gpu=False) as sess:
shape = (2, 3)
indices = array_ops.placeholder(shape=shape, dtype=np.int32)
data = np.zeros(shape + (5,))
partitions = data_flow_ops.dynamic_partition(
data, indices, num_partitions=7)
for i in range(2):
for j in range(3):
bad = np.zeros(shape, dtype=np.int32)
bad[i, j] = 17
with self.assertRaisesOpError(
r"partitions\[%d,%d\] = 17 is not in \[0, 7\)" % (i, j)):
sess.run(partitions, feed_dict={indices: bad})
@test_util.run_deprecated_v1
def testErrorWrongDimsIndices(self):
data = constant_op.constant([[0], [1], [2]])
indices = constant_op.constant([[0], [0]])
with self.assertRaises(ValueError):
data_flow_ops.dynamic_partition(data, indices, num_partitions=4)
# see https://github.com/tensorflow/tensorflow/issues/17106
def testCUBBug(self):
x = constant_op.constant(np.random.randn(3072))
inds = [0] * 189 + [1] * 184 + [2] * 184 + [3] * 191 + [4] * 192 + [
5
] * 195 + [6] * 195
inds += [7] * 195 + [8] * 188 + [9] * 195 + [10] * 188 + [11] * 202 + [
12
] * 194
inds += [13] * 194 + [14] * 194 + [15] * 192
self.assertEqual(len(inds), x.shape[0])
partitioned = data_flow_ops.dynamic_partition(x, inds, 16)
with self.cached_session():
res = self.evaluate(partitioned)
self.assertEqual(res[-1].shape[0], 192)
# see https://github.com/tensorflow/tensorflow/issues/42500
def testMultiGPU(self):
device_list = config.list_logical_devices("GPU")
results = []
for device in device_list:
with ops.device(device.name):
data = constant_op.constant(np.zeros((1000,)))
partitions = constant_op.constant(np.arange(1000, dtype=np.int32) % 10)
result = data_flow_ops.dynamic_partition(data, partitions, 10)
results.append(self.evaluate(result))
if device_list:
self.assertAllEqual(results, np.zeros((len(device_list), 10, 100)))
if __name__ == "__main__":
test.main()
tensorflow/python/kernel_tests/data_structures/dynamic_stitch_op_test.py
+400
View File
@@ -0,0 +1,400 @@
# Copyright 2015 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 the dynamic stitch and parallel dynamic stitch operations in data_flow_ops."""
import numpy as np
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import math_ops
import tensorflow.python.ops.data_flow_grad # pylint: disable=unused-import
from tensorflow.python.platform import test
class DynamicStitchTestBase(object):
def __init__(self, stitch_op):
self.stitch_op = stitch_op
def testScalar(self):
with test_util.use_gpu():
indices = [constant_op.constant(0), constant_op.constant(1)]
data = [constant_op.constant(40), constant_op.constant(60)]
for step in -1, 1:
stitched_t = self.stitch_op(indices[::step], data)
stitched_val = self.evaluate(stitched_t)
self.assertAllEqual([40, 60][::step], stitched_val)
# Dimension 0 is max(flatten(indices))+1.
self.assertEqual([2], stitched_t.get_shape().as_list())
@test_util.run_deprecated_v1
def testShapeInferenceForScalarWithNonConstantIndices(self):
with test_util.use_gpu():
indices = [
array_ops.placeholder(dtype=dtypes.int32),
constant_op.constant(1)
]
data = [constant_op.constant(40), constant_op.constant(60)]
for step in -1, 1:
stitched_t = self.stitch_op(indices[::step], data)
# Dimension 0 is max(flatten(indices))+1, but the first indices input is
# not a constant tensor, so we can only infer it as a vector of unknown
# length.
self.assertEqual([None], stitched_t.get_shape().as_list())
@test_util.disable_tfrt("b/169901260")
def testSimpleOneDimensional(self):
# Test various datatypes in the simple case to ensure that the op was
# registered under those types.
dtypes_to_test = [
dtypes.float32,
dtypes.float16,
dtypes.bfloat16,
dtypes.qint8,
dtypes.quint8,
dtypes.qint32,
]
for dtype in dtypes_to_test:
indices = [
constant_op.constant([0, 4, 7]),
constant_op.constant([1, 6, 2, 3, 5])
]
data = [
math_ops.cast(constant_op.constant([0, 40, 70]), dtype=dtype),
math_ops.cast(
constant_op.constant([10, 60, 20, 30, 50]), dtype=dtype)
]
stitched_t = self.stitch_op(indices, data)
stitched_val = self.evaluate(stitched_t)
self.assertAllEqual([0, 10, 20, 30, 40, 50, 60, 70], stitched_val)
# Dimension 0 is max(flatten(indices))+1.
self.assertEqual([8], stitched_t.get_shape().as_list())
def testOneListOneDimensional(self):
indices = [constant_op.constant([1, 6, 2, 3, 5, 0, 4, 7])]
data = [constant_op.constant([10, 60, 20, 30, 50, 0, 40, 70])]
stitched_t = self.stitch_op(indices, data)
stitched_val = self.evaluate(stitched_t)
self.assertAllEqual([0, 10, 20, 30, 40, 50, 60, 70], stitched_val)
# Dimension 0 is max(flatten(indices))+1.
self.assertEqual([8], stitched_t.get_shape().as_list())
def testSimpleTwoDimensional(self):
indices = [
constant_op.constant([0, 4, 7]),
constant_op.constant([1, 6]),
constant_op.constant([2, 3, 5])
]
data = [
constant_op.constant([[0, 1], [40, 41], [70, 71]]),
constant_op.constant([[10, 11], [60, 61]]),
constant_op.constant([[20, 21], [30, 31], [50, 51]])
]
stitched_t = self.stitch_op(indices, data)
stitched_val = self.evaluate(stitched_t)
self.assertAllEqual([[0, 1], [10, 11], [20, 21], [30, 31], [40, 41],
[50, 51], [60, 61], [70, 71]], stitched_val)
# Dimension 0 is max(flatten(indices))+1.
self.assertEqual([8, 2], stitched_t.get_shape().as_list())
def testZeroSizeTensor(self):
indices = [
constant_op.constant([0, 4, 7]),
constant_op.constant([1, 6]),
constant_op.constant([2, 3, 5]),
array_ops.zeros([0], dtype=dtypes.int32)
]
data = [
constant_op.constant([[0, 1], [40, 41], [70, 71]]),
constant_op.constant([[10, 11], [60, 61]]),
constant_op.constant([[20, 21], [30, 31], [50, 51]]),
array_ops.zeros([0, 2], dtype=dtypes.int32)
]
stitched_t = self.stitch_op(indices, data)
stitched_val = self.evaluate(stitched_t)
self.assertAllEqual([[0, 1], [10, 11], [20, 21], [30, 31], [40, 41],
[50, 51], [60, 61], [70, 71]], stitched_val)
# Dimension 0 is max(flatten(indices))+1.
self.assertEqual([8, 2], stitched_t.get_shape().as_list())
def testAllZeroSizeTensor(self):
indices = [
array_ops.zeros([0], dtype=dtypes.int32),
array_ops.zeros([0], dtype=dtypes.int32)
]
data = [
array_ops.zeros([0, 2], dtype=dtypes.int32),
array_ops.zeros([0, 2], dtype=dtypes.int32)
]
stitched_t = self.stitch_op(indices, data)
stitched_val = self.evaluate(stitched_t)
self.assertAllEqual(np.zeros((0, 2)), stitched_val)
self.assertEqual([0, 2], stitched_t.get_shape().as_list())
@test_util.run_deprecated_v1
def testHigherRank(self):
indices = [
constant_op.constant(6),
constant_op.constant([4, 1]),
constant_op.constant([[5, 2], [0, 3]])
]
data = [
constant_op.constant([61., 62.]),
constant_op.constant([[41., 42.], [11., 12.]]),
constant_op.constant([[[51., 52.], [21., 22.]],
[[1., 2.], [31., 32.]]])
]
stitched_t = self.stitch_op(indices, data)
stitched_val = self.evaluate(stitched_t)
correct = 10. * np.arange(7)[:, None] + [1., 2.]
self.assertAllEqual(correct, stitched_val)
self.assertEqual([7, 2], stitched_t.get_shape().as_list())
# Test gradients
stitched_grad = 7. * stitched_val
grads = gradients_impl.gradients(stitched_t, indices + data,
stitched_grad)
self.assertEqual(grads[:3], [None] * 3) # Indices have no gradients
for datum, grad in zip(data, self.evaluate(grads[3:])):
self.assertAllEqual(7. * self.evaluate(datum), grad)
@test_util.run_deprecated_v1
def testErrorIndicesMultiDimensional(self):
indices = [
constant_op.constant([0, 4, 7]),
constant_op.constant([[1, 6, 2, 3, 5]])
]
data = [
constant_op.constant([[0, 40, 70]]),
constant_op.constant([10, 60, 20, 30, 50])
]
with self.assertRaises(ValueError):
self.stitch_op(indices, data)
@test_util.run_deprecated_v1
def testErrorDataNumDimsMismatch(self):
indices = [
constant_op.constant([0, 4, 7]),
constant_op.constant([1, 6, 2, 3, 5])
]
data = [
constant_op.constant([0, 40, 70]),
constant_op.constant([[10, 60, 20, 30, 50]])
]
with self.assertRaises(ValueError):
self.stitch_op(indices, data)
@test_util.run_deprecated_v1
def testErrorDataDimSizeMismatch(self):
indices = [
constant_op.constant([0, 4, 5]),
constant_op.constant([1, 6, 2, 3])
]
data = [
constant_op.constant([[0], [40], [70]]),
constant_op.constant([[10, 11], [60, 61], [20, 21], [30, 31]])
]
with self.assertRaises(ValueError):
self.stitch_op(indices, data)
@test_util.run_deprecated_v1
def testErrorDataAndIndicesSizeMismatch(self):
indices = [
constant_op.constant([0, 4, 7]),
constant_op.constant([1, 6, 2, 3, 5])
]
data = [
constant_op.constant([0, 40, 70]),
constant_op.constant([10, 60, 20, 30])
]
with self.assertRaises(ValueError):
self.stitch_op(indices, data)
def testOutOfBoundsIndexRaisesInvalidArgument(self):
with self.assertRaisesRegex(errors.InvalidArgumentError, "out of range"):
indices = [[-1000], [405], [519], [758], [1015]]
data = [
[110.27793884277344],
[120.29475402832031],
[157.2418212890625],
[157.2626953125],
[188.45382690429688],
]
self.evaluate(self.stitch_op(indices, data))
@test_util.run_deprecated_v1
def testGradientWithDuplicateIndicesSingleGroup(self):
indices = [constant_op.constant([0, 0])]
data = [constant_op.constant([1.0, 2.0], dtype=dtypes.float32)]
stitched_t = self.stitch_op(indices, data)
stitched_grad = constant_op.constant([3.0], dtype=dtypes.float32)
grads = gradients_impl.gradients(stitched_t, indices + data, stitched_grad)
self.assertIsNone(grads[0]) # Indices have no gradient
self.assertAllEqual([0.0, 3.0], self.evaluate(grads[1]))
@test_util.run_deprecated_v1
def testGradientWithDuplicateIndicesMultiGroup(self):
indices = [constant_op.constant([0, 1]), constant_op.constant([1, 2])]
data = [
constant_op.constant([10.0, 20.0], dtype=dtypes.float32),
constant_op.constant([30.0, 40.0], dtype=dtypes.float32),
]
stitched_t = self.stitch_op(indices, data)
stitched_grad = constant_op.constant([2.0, 3.0, 4.0], dtype=dtypes.float32)
grads = gradients_impl.gradients(stitched_t, indices + data, stitched_grad)
self.assertEqual(grads[:2], [None] * 2) # Indices have no gradient
self.assertAllEqual([2.0, 0.0], self.evaluate(grads[2]))
self.assertAllEqual([3.0, 4.0], self.evaluate(grads[3]))
@test_util.run_deprecated_v1
def testGradientWithDuplicateIndicesHigherRank(self):
indices = [constant_op.constant([0, 0])]
data = [
constant_op.constant(
[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]],
dtype=dtypes.float32,
)
]
stitched_t = self.stitch_op(indices, data)
stitched_grad = constant_op.constant(
[[[10.0, 20.0], [30.0, 40.0]]], dtype=dtypes.float32
)
grads = gradients_impl.gradients(stitched_t, indices + data, stitched_grad)
self.assertIsNone(grads[0]) # Indices have no gradient
expected_grad = [[[0.0, 0.0], [0.0, 0.0]], [[10.0, 20.0], [30.0, 40.0]]]
self.assertAllEqual(expected_grad, self.evaluate(grads[1]))
@test_util.run_in_graph_and_eager_modes
def testMismatchedDataAndIndexListSizes(self):
indices = [
constant_op.constant([2]),
constant_op.constant([1]),
constant_op.constant([0]),
constant_op.constant([3]),
]
data = [
constant_op.constant([1.0]),
constant_op.constant([2.0]),
constant_op.constant([3.0]),
constant_op.constant([4.0]),
]
with self.assertRaisesRegex(
(ValueError, errors.InvalidArgumentError),
"expected inputs .* do not match|List argument .* must match",
):
self.evaluate(self.stitch_op(indices[0:2], data))
with self.assertRaisesRegex(
(ValueError, errors.InvalidArgumentError),
"expected inputs .* do not match|List argument .* must match",
):
self.evaluate(self.stitch_op(indices, data[0:2]))
class DynamicStitchTest(DynamicStitchTestBase, test.TestCase):
def __init__(self, *test_case_args):
test.TestCase.__init__(self, *test_case_args)
DynamicStitchTestBase.__init__(self, data_flow_ops.dynamic_stitch)
class ParallelDynamicStitchTest(DynamicStitchTestBase, test.TestCase):
def __init__(self, *test_case_args):
test.TestCase.__init__(self, *test_case_args)
DynamicStitchTestBase.__init__(self, data_flow_ops.parallel_dynamic_stitch)
def testScalar(self):
with test_util.use_gpu():
indices = [constant_op.constant(0), constant_op.constant(1)]
data = [constant_op.constant(40.0), constant_op.constant(60.0)]
for step in -1, 1:
stitched_t = self.stitch_op(indices[::step], data)
stitched_val = self.evaluate(stitched_t)
self.assertAllEqual([40.0, 60.0][::step], stitched_val)
# Dimension 0 is max(flatten(indices))+1.
self.assertEqual([2], stitched_t.get_shape().as_list())
@test_util.run_deprecated_v1
def testHigherRank(self):
indices = [
constant_op.constant(6),
constant_op.constant([4, 1]),
constant_op.constant([[5, 2], [0, 3]])
]
data = [
constant_op.constant([61, 62], dtype=dtypes.float32),
constant_op.constant([[41, 42], [11, 12]], dtype=dtypes.float32),
constant_op.constant(
[[[51, 52], [21, 22]], [[1, 2], [31, 32]]], dtype=dtypes.float32)
]
stitched_t = self.stitch_op(indices, data)
stitched_val = self.evaluate(stitched_t)
correct = 10 * np.arange(7)[:, None] + [1.0, 2.0]
self.assertAllEqual(correct, stitched_val)
self.assertEqual([7, 2], stitched_t.get_shape().as_list())
# Test gradients
stitched_grad = 7 * stitched_val
grads = gradients_impl.gradients(stitched_t, indices + data,
stitched_grad)
self.assertEqual(grads[:3], [None] * 3) # Indices have no gradients
for datum, grad in zip(data, self.evaluate(grads[3:])):
self.assertAllEqual(7.0 * self.evaluate(datum), grad)
# GPU version unit tests
def testScalarGPU(self):
indices = [constant_op.constant(0), constant_op.constant(1)]
data = [constant_op.constant(40.0), constant_op.constant(60.0)]
for step in -1, 1:
stitched_t = self.stitch_op(indices[::step], data)
stitched_val = self.evaluate(stitched_t)
self.assertAllEqual([40.0, 60.0][::step], stitched_val)
# Dimension 0 is max(flatten(indices))+1.
self.assertEqual([2], stitched_t.get_shape().as_list())
@test_util.run_deprecated_v1
def testHigherRankGPU(self):
indices = [
constant_op.constant(6),
constant_op.constant([4, 1]),
constant_op.constant([[5, 2], [0, 3]])
]
data = [
constant_op.constant([61, 62], dtype=dtypes.float32),
constant_op.constant([[41, 42], [11, 12]], dtype=dtypes.float32),
constant_op.constant(
[[[51, 52], [21, 22]], [[1, 2], [31, 32]]], dtype=dtypes.float32)
]
stitched_t = self.stitch_op(indices, data)
stitched_val = self.evaluate(stitched_t)
correct = 10 * np.arange(7)[:, None] + [1.0, 2.0]
self.assertAllEqual(correct, stitched_val)
self.assertEqual([7, 2], stitched_t.get_shape().as_list())
# Test gradients
stitched_grad = 7 * stitched_val
grads = gradients_impl.gradients(stitched_t, indices + data,
stitched_grad)
self.assertEqual(grads[:3], [None] * 3) # Indices have no gradients
for datum, grad in zip(data, self.evaluate(grads[3:])):
self.assertAllEqual(7.0 * self.evaluate(datum), grad)
if __name__ == "__main__":
test.main()
tensorflow/python/kernel_tests/data_structures/fifo_queue_test.py
+1785
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/kernel_tests/data_structures/list_ops_test.py
+1886
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/kernel_tests/data_structures/listdiff_op_test.py
+134
View File
@@ -0,0 +1,134 @@
# Copyright 2015 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 tensorflow.kernels.listdiff_op."""
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
from tensorflow.python.util import compat
_TYPES = [
dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64, dtypes.string
]
class ListDiffTest(test.TestCase):
def _testListDiff(self, x, y, out, idx):
for dtype in _TYPES:
if dtype == dtypes.string:
x = [compat.as_bytes(str(a)) for a in x]
y = [compat.as_bytes(str(a)) for a in y]
out = [compat.as_bytes(str(a)) for a in out]
for diff_func in [array_ops.setdiff1d]:
for index_dtype in [dtypes.int32, dtypes.int64]:
with self.cached_session() as sess:
x_tensor = ops.convert_to_tensor(x, dtype=dtype)
y_tensor = ops.convert_to_tensor(y, dtype=dtype)
out_tensor, idx_tensor = diff_func(x_tensor, y_tensor,
index_dtype=index_dtype)
tf_out, tf_idx = self.evaluate([out_tensor, idx_tensor])
self.assertAllEqual(tf_out, out)
self.assertAllEqual(tf_idx, idx)
self.assertEqual(1, out_tensor.get_shape().ndims)
self.assertEqual(1, idx_tensor.get_shape().ndims)
def testBasic1(self):
x = [1, 2, 3, 4]
y = [1, 2]
out = [3, 4]
idx = [2, 3]
self._testListDiff(x, y, out, idx)
def testBasic2(self):
x = [1, 2, 3, 4]
y = [2]
out = [1, 3, 4]
idx = [0, 2, 3]
self._testListDiff(x, y, out, idx)
def testBasic3(self):
x = [1, 4, 3, 2]
y = [4, 2]
out = [1, 3]
idx = [0, 2]
self._testListDiff(x, y, out, idx)
def testDuplicates(self):
x = [1, 2, 4, 3, 2, 3, 3, 1]
y = [4, 2]
out = [1, 3, 3, 3, 1]
idx = [0, 3, 5, 6, 7]
self._testListDiff(x, y, out, idx)
def testRandom(self):
num_random_tests = 10
int_low = -7
int_high = 8
max_size = 50
for _ in range(num_random_tests):
x_size = np.random.randint(max_size + 1)
x = np.random.randint(int_low, int_high, size=x_size)
y_size = np.random.randint(max_size + 1)
y = np.random.randint(int_low, int_high, size=y_size)
out_idx = [(entry, pos) for pos, entry in enumerate(x) if entry not in y]
if out_idx:
out, idx = map(list, zip(*out_idx))
else:
out = []
idx = []
self._testListDiff(list(x), list(y), out, idx)
def testFullyOverlapping(self):
x = [1, 2, 3, 4]
y = [1, 2, 3, 4]
out = []
idx = []
self._testListDiff(x, y, out, idx)
def testNonOverlapping(self):
x = [1, 2, 3, 4]
y = [5, 6]
out = x
idx = np.arange(len(x))
self._testListDiff(x, y, out, idx)
def testEmptyX(self):
x = []
y = [1, 2]
out = []
idx = []
self._testListDiff(x, y, out, idx)
def testEmptyY(self):
x = [1, 2, 3, 4]
y = []
out = x
idx = np.arange(len(x))
self._testListDiff(x, y, out, idx)
def testEmptyXY(self):
x = []
y = []
out = []
idx = []
self._testListDiff(x, y, out, idx)
if __name__ == "__main__":
test.main()
tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py
+4447
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/kernel_tests/data_structures/map_ops_test.py
+461
View File
@@ -0,0 +1,461 @@
# Copyright 2020 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 TensorMap ops."""
from absl.testing import parameterized
from tensorflow.python.eager import backprop
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import cond
from tensorflow.python.ops import map_ops
from tensorflow.python.ops import sort_ops
from tensorflow.python.platform import test
@test_util.run_all_in_graph_and_eager_modes
class MapOpsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
def testEmptyTensorMapSize(self):
m = map_ops.empty_tensor_map()
s = map_ops.tensor_map_size(m)
self.assertAllEqual(s, 0)
def testTensorMapInsert(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
v = constant_op.constant(2.0)
m = map_ops.tensor_map_insert(m, k, v)
s = map_ops.tensor_map_size(m)
self.assertAllEqual(s, 1)
def testTensorMapLookup(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
v = constant_op.constant(2.0)
m = map_ops.tensor_map_insert(m, k, v)
l = map_ops.tensor_map_lookup(m, k, dtypes.float32)
self.assertAllClose(l, v)
def testTensorMapLookupFromEmptyMapFails(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
with self.assertRaisesRegex(errors.InvalidArgumentError,
"Trying to lookup non-existent key. *"):
l = map_ops.tensor_map_lookup(m, k, dtypes.float32)
self.evaluate(l)
def testTensorMapLookupMissingKeyFails(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(11.0)
m = map_ops.tensor_map_insert(m, k, v)
with self.assertRaisesRegex(errors.InvalidArgumentError,
"Trying to lookup non-existent key. *"):
l = map_ops.tensor_map_lookup(m, k2, dtypes.float32)
self.evaluate(l)
def testTensorMapLookupDtypeMismatch(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
v = constant_op.constant(2.0, dtype=dtypes.float32)
m = map_ops.tensor_map_insert(m, k, v)
with self.assertRaisesRegex(
errors.InvalidArgumentError, "Key does not match requested dtype.*"
):
l = map_ops.tensor_map_lookup(m, k, dtypes.int32)
self.evaluate(l)
def testTensorMapErase(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
v = constant_op.constant(2.0)
m = map_ops.tensor_map_insert(m, k, v)
s = map_ops.tensor_map_size(m)
self.assertAllEqual(s, 1)
m = map_ops.tensor_map_erase(m, k, v.dtype)
s = map_ops.tensor_map_size(m)
self.assertAllEqual(s, 0)
def testTensorMapEraseFromEmptyMapFails(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
with self.assertRaisesRegex(errors.InvalidArgumentError,
"Trying to erase non-existent item. *"):
m = map_ops.tensor_map_erase(m, k, dtypes.float32)
self.evaluate(m)
def testTensorMapEraseMissingKeyFails(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(2.0)
m = map_ops.tensor_map_insert(m, k2, v)
with self.assertRaisesRegex(errors.InvalidArgumentError,
"Trying to erase non-existent item. *"):
m = map_ops.tensor_map_erase(m, k, dtypes.float32)
self.evaluate(m)
def testTensorMapHasKey(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(2.0)
m = map_ops.tensor_map_insert(m, k, v)
# Check has key.
b = map_ops.tensor_map_has_key(m, k)
b2 = map_ops.tensor_map_has_key(m, k2)
self.assertAllEqual(b, True)
self.assertAllEqual(b2, False)
def testIfHasKeyLookup(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(2.0)
m = map_ops.tensor_map_insert(m, k, v)
default_value = array_ops.zeros_like(v)
l = cond.cond(
map_ops.tensor_map_has_key(m, k),
lambda: map_ops.tensor_map_lookup(m, k, dtypes.float32),
lambda: default_value)
l2 = cond.cond(
map_ops.tensor_map_has_key(m, k2),
lambda: map_ops.tensor_map_lookup(m, k, dtypes.float32),
lambda: default_value)
self.assertAllClose(l, v)
self.assertAllClose(l2, default_value)
def testStackKeys(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
k3 = constant_op.constant(3.0)
v = constant_op.constant(21.0)
v2 = constant_op.constant(22.0)
v3 = constant_op.constant(23.0)
m = map_ops.tensor_map_insert(m, k, v)
m = map_ops.tensor_map_insert(m, k2, v2)
keys = map_ops.tensor_map_stack_keys(m, k.dtype)
expected = constant_op.constant([1.0, 2.0])
self.assertAllClose(array_ops.shape(keys), array_ops.shape(expected))
self.assertAllClose(sort_ops.sort(keys), expected)
m = map_ops.tensor_map_insert(m, k3, v3)
keys = map_ops.tensor_map_stack_keys(m, k.dtype)
expected = constant_op.constant([1.0, 2.0, 3.0])
self.assertAllClose(array_ops.shape(keys), array_ops.shape(expected))
self.assertAllClose(sort_ops.sort(keys), expected)
def testStackKeysEmptyMapFails(self):
m = map_ops.empty_tensor_map()
with self.assertRaisesRegex(
errors.InvalidArgumentError, "TensorMapStackKeys cannot be called "
"on empty map."):
keys = map_ops.tensor_map_stack_keys(m, dtypes.float32)
self.evaluate(keys)
def testStackKeysIncorrectDtypeFails(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant("key_with_wrong_dtype")
v = constant_op.constant(2.0)
m = map_ops.tensor_map_insert(m, k, v)
simple = "Key does not match requested dtype."
with self.assertRaisesRegex(errors.InvalidArgumentError, simple):
keys = map_ops.tensor_map_stack_keys(m, dtypes.float32)
self.evaluate(keys)
def testStackKeysIncorrectShapeFails(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant([1.0, 11.0])
v = constant_op.constant(2.0)
v2 = constant_op.constant(22.0)
m = map_ops.tensor_map_insert(m, k, v)
m = map_ops.tensor_map_insert(m, k2, v2)
with self.assertRaisesRegex(errors.InvalidArgumentError,
"Keys must all have the same shape."):
keys = map_ops.tensor_map_stack_keys(m, dtypes.float32)
self.evaluate(keys)
def testInsertLookupGrad(self):
with backprop.GradientTape() as tape:
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
v = constant_op.constant(11.0)
tape.watch(v)
m = map_ops.tensor_map_insert(m, k, v)
l = map_ops.tensor_map_lookup(m, k, dtypes.float32)
l *= 5
g = tape.gradient(l, v)
self.assertAllEqual(g, 5)
def testMultipleInsertLookupGrad(self):
with backprop.GradientTape(persistent=True) as tape:
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
k3 = constant_op.constant(3.0)
v = constant_op.constant(11.0)
v2 = constant_op.constant(12.0)
v3 = constant_op.constant(13.0)
tape.watch(v)
tape.watch(v2)
tape.watch(v3)
m = map_ops.tensor_map_insert(m, k, v)
m = map_ops.tensor_map_insert(m, k2, v2)
m = map_ops.tensor_map_insert(m, k3, v3)
l = map_ops.tensor_map_lookup(m, k, v.dtype)
l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype)
l3 = map_ops.tensor_map_lookup(m, k3, v3.dtype)
g = tape.gradient(l * 5, v)
g2 = tape.gradient(l2 * 6, v2)
g3 = tape.gradient(l3 * 7, v3)
self.assertAllEqual(g, 5)
self.assertAllEqual(g2, 6)
self.assertAllEqual(g3, 7)
del tape
def testInsertLookupComposeGrad(self):
with backprop.GradientTape() as tape:
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(11.0)
tape.watch(v)
m = map_ops.tensor_map_insert(m, k, v)
l = map_ops.tensor_map_lookup(m, k, v.dtype)
m = map_ops.tensor_map_insert(m, k2, l)
l2 = map_ops.tensor_map_lookup(m, k2, l.dtype)
g = tape.gradient(l2 * 5, v)
self.assertAllEqual(g, 5)
def testReplaceLookupGrad(self):
# Test using same key and different value.
with backprop.GradientTape(persistent=True) as tape:
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
v = constant_op.constant(11.0)
v2 = constant_op.constant(22.0)
tape.watch(v)
tape.watch(v2)
m = map_ops.tensor_map_insert(m, k, v)
l = map_ops.tensor_map_lookup(m, k, v.dtype)
self.assertAllClose(l, v)
g = tape.gradient(l * 5, v)
self.assertAllEqual(g, 5)
# Replace key and lookup.
m = map_ops.tensor_map_insert(m, k, v2)
l2 = map_ops.tensor_map_lookup(m, k, v2.dtype)
self.assertAllClose(l2, v2)
g2 = tape.gradient(l2 * 6, v)
self.assertAllClose(g2, array_ops.zeros_like(v))
g3 = tape.gradient(l2 * 7, v2)
self.assertAllClose(g3, 7)
del tape
def testDiffKeySameValueGrad(self):
with backprop.GradientTape(persistent=True) as tape:
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(11.0)
v = constant_op.constant(2.0)
v2 = constant_op.constant(22.0)
tape.watch(v)
tape.watch(v2)
m = map_ops.tensor_map_insert(m, k, v)
m = map_ops.tensor_map_insert(m, k2, v)
l = map_ops.tensor_map_lookup(m, k, v.dtype)
l2 = map_ops.tensor_map_lookup(m, k2, v.dtype)
g = tape.gradient(l + l2, v)
self.assertAllEqual(g, 2)
m = map_ops.tensor_map_insert(m, k2, v2)
l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype)
g2 = tape.gradient(l + l2, v2)
self.assertAllEqual(g2, 1)
del tape
def testLookupAddGrad(self):
with backprop.GradientTape(persistent=True) as tape:
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(11.0)
v2 = constant_op.constant(22.0)
tape.watch(v)
tape.watch(v2)
m = map_ops.empty_tensor_map()
m = map_ops.tensor_map_insert(m, k, v)
m = map_ops.tensor_map_insert(m, k2, v2)
l1 = map_ops.tensor_map_lookup(m, k, v.dtype)
l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype)
g = tape.gradient(l1 + l2, [l1, l2])
self.assertAllClose(g, [1, 1])
g2 = tape.gradient(l1 + l2, [v, v2])
self.assertAllClose(g2, [1, 1])
g3 = tape.gradient(l1 + l2 * 4, v2)
self.assertAllEqual(g3, 4)
del tape
def testLookupMultiplyGrad(self):
with backprop.GradientTape(persistent=True) as tape:
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(11.0)
v2 = constant_op.constant(22.0)
tape.watch(v)
tape.watch(v2)
m = map_ops.empty_tensor_map()
m = map_ops.tensor_map_insert(m, k, v)
m = map_ops.tensor_map_insert(m, k2, v2)
l1 = map_ops.tensor_map_lookup(m, k, v.dtype)
l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype)
g = tape.gradient(l1 * l2, [v, v2])
self.assertAllClose(g, [v2, v])
g2 = tape.gradient(l1 * l1, v)
self.assertAllClose(g2, 2 * v)
del tape
def testEraseFirstInsertGrad(self):
with backprop.GradientTape(persistent=True) as tape:
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(11.0)
v2 = constant_op.constant(22.0)
tape.watch(v)
tape.watch(v2)
m = map_ops.tensor_map_insert(m, k, v)
l = map_ops.tensor_map_lookup(m, k, v.dtype)
m = map_ops.tensor_map_insert(m, k2, v2)
m = map_ops.tensor_map_erase(m, k, v.dtype)
l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype)
self.assertAllClose(l2, v2)
g = tape.gradient(l * 5, v)
self.assertAllEqual(g, 5)
g2 = tape.gradient(l2 * 6, v2)
self.assertAllEqual(g2, 6)
del tape
def testEraseSecondInsertGrad(self):
with backprop.GradientTape(persistent=True) as tape:
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(11.0)
v2 = constant_op.constant(22.0)
tape.watch(v)
tape.watch(v2)
m = map_ops.tensor_map_insert(m, k, v)
m = map_ops.tensor_map_insert(m, k2, v2)
m = map_ops.tensor_map_erase(m, k2, v2.dtype)
l = map_ops.tensor_map_lookup(m, k, v.dtype)
self.assertAllClose(l, v)
g = tape.gradient(l * 5, v)
self.assertAllEqual(g, 5)
del tape
def testEraseInsertComposedGrad(self):
with backprop.GradientTape(persistent=True) as tape:
m = map_ops.empty_tensor_map()
k = constant_op.constant(1.0)
k2 = constant_op.constant(2.0)
v = constant_op.constant(11.0)
v2 = constant_op.constant(22.0)
tape.watch(v)
tape.watch(v2)
m = map_ops.tensor_map_insert(m, k, v)
l = map_ops.tensor_map_lookup(m, k, v.dtype)
m = map_ops.tensor_map_erase(m, k, v.dtype)
m = map_ops.tensor_map_insert(m, k2, l)
l2 = map_ops.tensor_map_lookup(m, k2, l.dtype)
g = tape.gradient(l2 * 5, v)
self.assertAllEqual(g, 5)
del tape
def testStringKeyGrad(self):
with backprop.GradientTape(persistent=True) as tape:
m = map_ops.empty_tensor_map()
k = constant_op.constant("key")
k2 = constant_op.constant("key2")
v = constant_op.constant(2.0)
v2 = constant_op.constant(22.0)
tape.watch(v)
tape.watch(v2)
m = map_ops.tensor_map_insert(m, k, v)
m = map_ops.tensor_map_insert(m, k2, v2)
s = map_ops.tensor_map_size(m)
self.assertAllEqual(s, 2)
# Test lookup and gradient.
l = map_ops.tensor_map_lookup(m, k, v.dtype)
self.assertAllClose(l, v)
self.assertAllClose(tape.gradient(l * 5, v), 5)
# Test replace and gradient.
m = map_ops.tensor_map_insert(m, k, v2)
l2 = map_ops.tensor_map_lookup(m, k, v2.dtype)
self.assertAllClose(l2, v2)
g = tape.gradient(l2 * 6, v2)
self.assertAllEqual(g, 6)
# Test erase, has key, and gradient.
m = map_ops.tensor_map_erase(m, k, v2.dtype)
s = map_ops.tensor_map_size(m)
self.assertAllEqual(s, 1)
h = map_ops.tensor_map_has_key(m, k)
self.assertAllEqual(h, False)
l = map_ops.tensor_map_lookup(m, k2, v2.dtype)
g2 = tape.gradient(l * 6, v2)
self.assertAllEqual(g2, 6)
del tape
def testStringKeyValue(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant("key")
v = constant_op.constant("value")
k2 = constant_op.constant(1.0)
v2 = constant_op.constant(2.0)
# Test insert and lookup on string key-value pair.
m = map_ops.tensor_map_insert(m, k, v)
m = map_ops.tensor_map_insert(m, k2, v2)
l = map_ops.tensor_map_lookup(m, k, v.dtype)
self.assertAllEqual(l, v)
# Test lookup on float key-value pair.
l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype)
self.assertAllClose(l2, v2)
# Test erase and has key.
self.assertAllEqual(map_ops.tensor_map_has_key(m, k), True)
m = map_ops.tensor_map_erase(m, k, v.dtype)
self.assertAllEqual(map_ops.tensor_map_has_key(m, k), False)
self.assertAllEqual(map_ops.tensor_map_has_key(m, k2), True)
def testVectorValue(self):
m = map_ops.empty_tensor_map()
k = constant_op.constant([1.0, 2.0])
v = constant_op.constant([11.0, 22.0])
# Test insert and lookup.
m = map_ops.tensor_map_insert(m, k, v)
s = map_ops.tensor_map_size(m)
self.assertAllEqual(s, 1)
l = map_ops.tensor_map_lookup(m, k, v.dtype)
self.assertAllEqual(l, v)
# Test erase and has key.
m = map_ops.tensor_map_erase(m, k, v.dtype)
s = map_ops.tensor_map_size(m)
self.assertAllEqual(s, 0)
self.assertAllEqual(map_ops.tensor_map_has_key(m, k), False)
if __name__ == "__main__":
test.main()
tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py
+660
View File
@@ -0,0 +1,660 @@
# Copyright 2017 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.
# ==============================================================================
import queue
import threading
import numpy as np
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
TIMEOUT = 1
class MapStageTest(test.TestCase):
@test_util.run_deprecated_v1
def testSimple(self):
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
pi = array_ops.placeholder(dtypes.int64)
gi = array_ops.placeholder(dtypes.int64)
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.MapStagingArea([dtypes.float32])
stage = stager.put(pi, [v], [0])
k, y = stager.get(gi)
y = math_ops.reduce_max(math_ops.matmul(y, y))
g.finalize()
with self.session(graph=g) as sess:
sess.run(stage, feed_dict={x: -1, pi: 0})
for i in range(10):
_, yval = sess.run([stage, y], feed_dict={x: i, pi: i + 1, gi: i})
self.assertAllClose(4 * (i - 1) * (i - 1) * 128, yval, rtol=1e-4)
@test_util.run_deprecated_v1
def testMultiple(self):
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
pi = array_ops.placeholder(dtypes.int64)
gi = array_ops.placeholder(dtypes.int64)
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32])
stage = stager.put(pi, [x, v], [0, 1])
k, (z, y) = stager.get(gi)
y = math_ops.reduce_max(z * math_ops.matmul(y, y))
g.finalize()
with self.session(graph=g) as sess:
sess.run(stage, feed_dict={x: -1, pi: 0})
for i in range(10):
_, yval = sess.run([stage, y], feed_dict={x: i, pi: i + 1, gi: i})
self.assertAllClose(
4 * (i - 1) * (i - 1) * (i - 1) * 128, yval, rtol=1e-4)
@test_util.run_deprecated_v1
def testDictionary(self):
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
pi = array_ops.placeholder(dtypes.int64)
gi = array_ops.placeholder(dtypes.int64)
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32],
shapes=[[], [128, 128]],
names=['x', 'v'])
stage = stager.put(pi, {'x': x, 'v': v})
key, ret = stager.get(gi)
z = ret['x']
y = ret['v']
y = math_ops.reduce_max(z * math_ops.matmul(y, y))
g.finalize()
with self.session(graph=g) as sess:
sess.run(stage, feed_dict={x: -1, pi: 0})
for i in range(10):
_, yval = sess.run([stage, y], feed_dict={x: i, pi: i + 1, gi: i})
self.assertAllClose(
4 * (i - 1) * (i - 1) * (i - 1) * 128, yval, rtol=1e-4)
def testColocation(self):
gpu_dev = test.gpu_device_name()
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(gpu_dev):
stager = data_flow_ops.MapStagingArea([dtypes.float32])
y = stager.put(1, [v], [0])
expected_name = gpu_dev if 'gpu' not in gpu_dev else '/device:GPU:0'
self.assertEqual(y.device, expected_name)
with ops.device('/cpu:0'):
_, x = stager.get(1)
y = stager.peek(1)[0]
_, z = stager.get()
self.assertEqual(x[0].device, '/device:CPU:0')
self.assertEqual(y.device, '/device:CPU:0')
self.assertEqual(z[0].device, '/device:CPU:0')
g.finalize()
@test_util.run_deprecated_v1
def testPeek(self):
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.int32, name='x')
pi = array_ops.placeholder(dtypes.int64)
gi = array_ops.placeholder(dtypes.int64)
p = array_ops.placeholder(dtypes.int32, name='p')
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.MapStagingArea([
dtypes.int32,
], shapes=[[]])
stage = stager.put(pi, [x], [0])
peek = stager.peek(gi)
size = stager.size()
g.finalize()
n = 10
with self.session(graph=g) as sess:
for i in range(n):
sess.run(stage, feed_dict={x: i, pi: i})
for i in range(n):
self.assertEqual(sess.run(peek, feed_dict={gi: i})[0], i)
self.assertEqual(sess.run(size), 10)
@test_util.run_deprecated_v1
def testSizeAndClear(self):
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32, name='x')
pi = array_ops.placeholder(dtypes.int64)
gi = array_ops.placeholder(dtypes.int64)
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32],
shapes=[[], [128, 128]],
names=['x', 'v'])
stage = stager.put(pi, {'x': x, 'v': v})
size = stager.size()
clear = stager.clear()
g.finalize()
with self.session(graph=g) as sess:
sess.run(stage, feed_dict={x: -1, pi: 3})
self.assertEqual(sess.run(size), 1)
sess.run(stage, feed_dict={x: -1, pi: 1})
self.assertEqual(sess.run(size), 2)
sess.run(clear)
self.assertEqual(sess.run(size), 0)
@test_util.run_deprecated_v1
def testCapacity(self):
capacity = 3
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.int32, name='x')
pi = array_ops.placeholder(dtypes.int64, name='pi')
gi = array_ops.placeholder(dtypes.int64, name='gi')
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.MapStagingArea([
dtypes.int32,
],
capacity=capacity,
shapes=[[]])
stage = stager.put(pi, [x], [0])
get = stager.get()
size = stager.size()
g.finalize()
value_queue = queue.Queue()
n = 8
with self.session(graph=g) as sess:
# Stage data in a separate thread which will block when it hits the
# staging area's capacity and thus not fill the value_queue with n tokens
def thread_run():
for i in range(n):
sess.run(stage, feed_dict={x: i, pi: i})
value_queue.put(0)
t = threading.Thread(target=thread_run)
t.daemon = True
t.start()
# Get tokens from the value_queue until a timeout occurs
try:
for i in range(n):
value_queue.get(timeout=TIMEOUT)
except queue.Empty:
pass
# Should've timed out on the iteration 'capacity'
if not i == capacity:
self.fail("Expected to timeout on iteration '{}' "
"but instead timed out on iteration '{}' "
"Staging Area size is '{}' and configured "
"capacity is '{}'.".format(capacity, i, sess.run(size),
capacity))
# Should have capacity elements in the staging area
self.assertEqual(sess.run(size), capacity)
# Clear the staging area completely
for i in range(n):
sess.run(get)
self.assertEqual(sess.run(size), 0)
@test_util.run_deprecated_v1
def testMemoryLimit(self):
memory_limit = 512 * 1024 # 512K
chunk = 200 * 1024 # 256K
capacity = memory_limit // chunk
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.uint8, name='x')
pi = array_ops.placeholder(dtypes.int64, name='pi')
gi = array_ops.placeholder(dtypes.int64, name='gi')
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.MapStagingArea([dtypes.uint8],
memory_limit=memory_limit,
shapes=[[]])
stage = stager.put(pi, [x], [0])
get = stager.get()
size = stager.size()
g.finalize()
value_queue = queue.Queue()
n = 8
with self.session(graph=g) as sess:
# Stage data in a separate thread which will block when it hits the
# staging area's capacity and thus not fill the value_queue with n tokens
def thread_run():
for i in range(n):
data = np.full(chunk, i, dtype=np.uint8)
sess.run(stage, feed_dict={x: data, pi: i})
value_queue.put(0)
t = threading.Thread(target=thread_run)
t.daemon = True
t.start()
# Get tokens from the value_queue until a timeout occurs
try:
for i in range(n):
value_queue.get(timeout=TIMEOUT)
except queue.Empty:
pass
# Should've timed out on the iteration 'capacity'
if not i == capacity:
self.fail("Expected to timeout on iteration '{}' "
"but instead timed out on iteration '{}' "
"Staging Area size is '{}' and configured "
"capacity is '{}'.".format(capacity, i, sess.run(size),
capacity))
# Should have capacity elements in the staging area
self.assertEqual(sess.run(size), capacity)
# Clear the staging area completely
for i in range(n):
sess.run(get)
self.assertEqual(sess.run(size), 0)
@test_util.run_deprecated_v1
def testOrdering(self):
import random
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.int32, name='x')
pi = array_ops.placeholder(dtypes.int64, name='pi')
gi = array_ops.placeholder(dtypes.int64, name='gi')
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.MapStagingArea([
dtypes.int32,
],
shapes=[[]],
ordered=True)
stage = stager.put(pi, [x], [0])
get = stager.get()
size = stager.size()
g.finalize()
n = 10
with self.session(graph=g) as sess:
# Keys n-1..0
keys = list(reversed(range(n)))
for i in keys:
sess.run(stage, feed_dict={pi: i, x: i})
self.assertEqual(sess.run(size), n)
# Check that key, values come out in ascending order
for i, k in enumerate(reversed(keys)):
get_key, values = sess.run(get)
self.assertTrue(i == k == get_key == values)
self.assertEqual(sess.run(size), 0)
@test_util.run_deprecated_v1
def testPartialDictInsert(self):
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
f = array_ops.placeholder(dtypes.float32)
v = array_ops.placeholder(dtypes.float32)
pi = array_ops.placeholder(dtypes.int64)
gi = array_ops.placeholder(dtypes.int64)
with ops.device(test.gpu_device_name()):
# Test barrier with dictionary
stager = data_flow_ops.MapStagingArea(
[dtypes.float32, dtypes.float32, dtypes.float32],
names=['x', 'v', 'f'])
stage_xf = stager.put(pi, {'x': x, 'f': f})
stage_v = stager.put(pi, {'v': v})
key, ret = stager.get(gi)
size = stager.size()
isize = stager.incomplete_size()
g.finalize()
with self.session(graph=g) as sess:
# 0 complete and incomplete entries
self.assertEqual(sess.run([size, isize]), [0, 0])
# Stage key 0, x and f tuple entries
sess.run(stage_xf, feed_dict={pi: 0, x: 1, f: 2})
self.assertEqual(sess.run([size, isize]), [0, 1])
# Stage key 1, x and f tuple entries
sess.run(stage_xf, feed_dict={pi: 1, x: 1, f: 2})
self.assertEqual(sess.run([size, isize]), [0, 2])
# Now complete key 0 with tuple entry v
sess.run(stage_v, feed_dict={pi: 0, v: 1})
# 1 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [1, 1])
# We can now obtain tuple associated with key 0
self.assertEqual(
sess.run([key, ret], feed_dict={gi: 0}),
[0, {
'x': 1,
'f': 2,
'v': 1
}])
# 0 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [0, 1])
# Now complete key 1 with tuple entry v
sess.run(stage_v, feed_dict={pi: 1, v: 3})
# We can now obtain tuple associated with key 1
self.assertEqual(
sess.run([key, ret], feed_dict={gi: 1}),
[1, {
'x': 1,
'f': 2,
'v': 3
}])
@test_util.run_deprecated_v1
def testPartialIndexInsert(self):
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
f = array_ops.placeholder(dtypes.float32)
v = array_ops.placeholder(dtypes.float32)
pi = array_ops.placeholder(dtypes.int64)
gi = array_ops.placeholder(dtypes.int64)
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.MapStagingArea(
[dtypes.float32, dtypes.float32, dtypes.float32])
stage_xf = stager.put(pi, [x, f], [0, 2])
stage_v = stager.put(pi, [v], [1])
key, ret = stager.get(gi)
size = stager.size()
isize = stager.incomplete_size()
g.finalize()
with self.session(graph=g) as sess:
# 0 complete and incomplete entries
self.assertEqual(sess.run([size, isize]), [0, 0])
# Stage key 0, x and f tuple entries
sess.run(stage_xf, feed_dict={pi: 0, x: 1, f: 2})
self.assertEqual(sess.run([size, isize]), [0, 1])
# Stage key 1, x and f tuple entries
sess.run(stage_xf, feed_dict={pi: 1, x: 1, f: 2})
self.assertEqual(sess.run([size, isize]), [0, 2])
# Now complete key 0 with tuple entry v
sess.run(stage_v, feed_dict={pi: 0, v: 1})
# 1 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [1, 1])
# We can now obtain tuple associated with key 0
self.assertEqual(sess.run([key, ret], feed_dict={gi: 0}), [0, [1, 1, 2]])
# 0 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [0, 1])
# Now complete key 1 with tuple entry v
sess.run(stage_v, feed_dict={pi: 1, v: 3})
# We can now obtain tuple associated with key 1
self.assertEqual(sess.run([key, ret], feed_dict={gi: 1}), [1, [1, 3, 2]])
@test_util.run_deprecated_v1
def testPartialDictGetsAndPeeks(self):
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
f = array_ops.placeholder(dtypes.float32)
v = array_ops.placeholder(dtypes.float32)
pi = array_ops.placeholder(dtypes.int64)
pei = array_ops.placeholder(dtypes.int64)
gi = array_ops.placeholder(dtypes.int64)
with ops.device(test.gpu_device_name()):
# Test barrier with dictionary
stager = data_flow_ops.MapStagingArea(
[dtypes.float32, dtypes.float32, dtypes.float32],
names=['x', 'v', 'f'])
stage_xf = stager.put(pi, {'x': x, 'f': f})
stage_v = stager.put(pi, {'v': v})
peek_xf = stager.peek(pei, ['x', 'f'])
peek_v = stager.peek(pei, ['v'])
key_xf, get_xf = stager.get(gi, ['x', 'f'])
key_v, get_v = stager.get(gi, ['v'])
pop_key_xf, pop_xf = stager.get(indices=['x', 'f'])
pop_key_v, pop_v = stager.get(pi, ['v'])
size = stager.size()
isize = stager.incomplete_size()
g.finalize()
with self.session(graph=g) as sess:
# 0 complete and incomplete entries
self.assertEqual(sess.run([size, isize]), [0, 0])
# Stage key 0, x and f tuple entries
sess.run(stage_xf, feed_dict={pi: 0, x: 1, f: 2})
self.assertEqual(sess.run([size, isize]), [0, 1])
# Stage key 1, x and f tuple entries
sess.run(stage_xf, feed_dict={pi: 1, x: 1, f: 2})
self.assertEqual(sess.run([size, isize]), [0, 2])
# Now complete key 0 with tuple entry v
sess.run(stage_v, feed_dict={pi: 0, v: 1})
# 1 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [1, 1])
# We can now peek at 'x' and 'f' values associated with key 0
self.assertEqual(sess.run(peek_xf, feed_dict={pei: 0}), {'x': 1, 'f': 2})
# Peek at 'v' value associated with key 0
self.assertEqual(sess.run(peek_v, feed_dict={pei: 0}), {'v': 1})
# 1 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [1, 1])
# We can now obtain 'x' and 'f' values associated with key 0
self.assertEqual(
sess.run([key_xf, get_xf], feed_dict={gi: 0}), [0, {
'x': 1,
'f': 2
}])
# Still have 1 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [1, 1])
# We can no longer get 'x' and 'f' from key 0
with self.assertRaises(errors.InvalidArgumentError) as cm:
sess.run([key_xf, get_xf], feed_dict={gi: 0})
exc_str = ("Tensor at index '0' for key '0' " 'has already been removed.')
self.assertIn(exc_str, cm.exception.message)
# Obtain 'v' value associated with key 0
self.assertEqual(
sess.run([key_v, get_v], feed_dict={gi: 0}), [0, {
'v': 1
}])
# 0 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [0, 1])
# Now complete key 1 with tuple entry v
sess.run(stage_v, feed_dict={pi: 1, v: 1})
# 1 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [1, 0])
# Pop without key to obtain 'x' and 'f' values associated with key 1
self.assertEqual(sess.run([pop_key_xf, pop_xf]), [1, {'x': 1, 'f': 2}])
# still 1 complete and 1 incomplete entry
self.assertEqual(sess.run([size, isize]), [1, 0])
# We can now obtain 'x' and 'f' values associated with key 1
self.assertEqual(
sess.run([pop_key_v, pop_v], feed_dict={pi: 1}), [1, {
'v': 1
}])
# Nothing is left
self.assertEqual(sess.run([size, isize]), [0, 0])
@test_util.run_deprecated_v1
def testPartialIndexGets(self):
with ops.Graph().as_default() as g:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
f = array_ops.placeholder(dtypes.float32)
v = array_ops.placeholder(dtypes.float32)
pi = array_ops.placeholder(dtypes.int64)
pei = array_ops.placeholder(dtypes.int64)
gi = array_ops.placeholder(dtypes.int64)
with ops.device(test.gpu_device_name()):
# Test again with partial index gets
stager = data_flow_ops.MapStagingArea(
[dtypes.float32, dtypes.float32, dtypes.float32])
stage_xvf = stager.put(pi, [x, v, f], [0, 1, 2])
key_xf, get_xf = stager.get(gi, [0, 2])
key_v, get_v = stager.get(gi, [1])
size = stager.size()
isize = stager.incomplete_size()
g.finalize()
with self.session(graph=g) as sess:
# Stage complete tuple
sess.run(stage_xvf, feed_dict={pi: 0, x: 1, f: 2, v: 3})
self.assertEqual(sess.run([size, isize]), [1, 0])
# Partial get using indices
self.assertEqual(
sess.run([key_xf, get_xf], feed_dict={gi: 0}), [0, [1, 2]])
# Still some of key 0 left
self.assertEqual(sess.run([size, isize]), [1, 0])
# Partial get of remaining index
self.assertEqual(sess.run([key_v, get_v], feed_dict={gi: 0}), [0, [3]])
# All gone
self.assertEqual(sess.run([size, isize]), [0, 0])
@test_util.run_deprecated_v1
def testNonScalarKeyOrderedMap(self):
with ops.Graph().as_default() as g:
x = array_ops.placeholder(dtypes.float32)
v = 2. * (array_ops.zeros([128, 128]) + x)
t = data_flow_ops.gen_data_flow_ops.ordered_map_stage(
key=constant_op.constant(value=[1], shape=(1, 3), dtype=dtypes.int64),
indices=np.array([[6]]),
values=[x, v],
dtypes=[dtypes.int64],
capacity=0,
memory_limit=0,
container='container1',
shared_name='',
name=None)
g.finalize()
with self.session(graph=g) as sess:
with self.assertRaisesRegex(errors.InvalidArgumentError,
'key must be an int64 scalar'):
sess.run(t, feed_dict={x: 1})
@test_util.run_deprecated_v1
def testNonScalarKeyUnorderedMap(self):
with ops.Graph().as_default() as g:
x = array_ops.placeholder(dtypes.float32)
v = 2. * (array_ops.zeros([128, 128]) + x)
t = data_flow_ops.gen_data_flow_ops.map_stage(
key=constant_op.constant(value=[1], shape=(1, 3), dtype=dtypes.int64),
indices=np.array([[6]]),
values=[x, v],
dtypes=[dtypes.int64],
capacity=0,
memory_limit=0,
container='container1',
shared_name='',
name=None)
g.finalize()
with self.session(graph=g) as sess:
with self.assertRaisesRegex(errors.InvalidArgumentError,
'key must be an int64 scalar'):
sess.run(t, feed_dict={x: 1})
def testNonScalarKeyMapPeek(self):
with self.assertRaisesRegex(
errors.InvalidArgumentError, 'key must be an int64 scalar'
):
v = data_flow_ops.gen_data_flow_ops.map_peek(
key=constant_op.constant(value=[1], shape=(1, 3), dtype=dtypes.int64),
indices=np.array([[6]]),
dtypes=[dtypes.int64],
capacity=0,
memory_limit=0,
container='container1',
shared_name='',
name=None,
)
self.evaluate(v)
def testNonScalarKeyMapUnStage(self):
with self.assertRaisesRegex(
errors.InvalidArgumentError, 'key must be an int64 scalar'
):
v = data_flow_ops.gen_data_flow_ops.map_unstage(
key=constant_op.constant(value=[1], shape=(1, 3), dtype=dtypes.int64),
indices=np.array([[6]]),
dtypes=[dtypes.int64],
capacity=0,
memory_limit=0,
container='container1',
shared_name='',
name=None,
)
self.evaluate(v)
if __name__ == '__main__':
test.main()
tensorflow/python/kernel_tests/data_structures/padding_fifo_queue_test.py
+1651
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/kernel_tests/data_structures/priority_queue_test.py
+354
View File
@@ -0,0 +1,354 @@
# 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 tensorflow.ops.data_flow_ops.PriorityQueue."""
import copy
import random
import threading
import numpy as np
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
import tensorflow.python.ops.nn_grad # pylint: disable=unused-import
from tensorflow.python.platform import test
@test_util.run_v1_only("PriorityQueue removed from v2")
class PriorityQueueTest(test.TestCase):
def testRoundTripInsertReadOnceSorts(self):
with self.cached_session() as sess:
q = data_flow_ops.PriorityQueue(2000, (dtypes.string, dtypes.string), (
(), ()))
elem = np.random.randint(-5, 5, size=100).astype(np.int64)
side_value_0 = np.random.rand(100).astype(bytes)
side_value_1 = np.random.rand(100).astype(bytes)
enq_list = [
q.enqueue((e, constant_op.constant(v0), constant_op.constant(v1)))
for e, v0, v1 in zip(elem, side_value_0, side_value_1)
]
for enq in enq_list:
enq.run()
deq = q.dequeue_many(100)
deq_elem, deq_value_0, deq_value_1 = self.evaluate(deq)
allowed = {}
missed = set()
for e, v0, v1 in zip(elem, side_value_0, side_value_1):
if e not in allowed:
allowed[e] = set()
allowed[e].add((v0, v1))
missed.add((v0, v1))
self.assertAllEqual(deq_elem, sorted(elem))
for e, dv0, dv1 in zip(deq_elem, deq_value_0, deq_value_1):
self.assertTrue((dv0, dv1) in allowed[e])
missed.remove((dv0, dv1))
self.assertEqual(missed, set())
def testRoundTripInsertMultiThreadedReadOnceSorts(self):
# We need each thread to keep its own device stack or the device scopes
# won't be properly nested.
ops.get_default_graph().switch_to_thread_local()
with self.cached_session() as sess:
q = data_flow_ops.PriorityQueue(2000, (dtypes.string, dtypes.string), (
(), ()))
elem = np.random.randint(-5, 5, size=100).astype(np.int64)
side_value_0 = np.random.rand(100).astype(bytes)
side_value_1 = np.random.rand(100).astype(bytes)
enqueue_ops = [
q.enqueue((e, constant_op.constant(v0), constant_op.constant(v1)))
for e, v0, v1 in zip(elem, side_value_0, side_value_1)
]
# Run one producer thread for each element in elems.
def enqueue(enqueue_op):
self.evaluate(enqueue_op)
dequeue_op = q.dequeue_many(100)
enqueue_threads = [
self.checkedThread(
target=enqueue, args=(op,)) for op in enqueue_ops
]
for t in enqueue_threads:
t.start()
deq_elem, deq_value_0, deq_value_1 = self.evaluate(dequeue_op)
for t in enqueue_threads:
t.join()
allowed = {}
missed = set()
for e, v0, v1 in zip(elem, side_value_0, side_value_1):
if e not in allowed:
allowed[e] = set()
allowed[e].add((v0, v1))
missed.add((v0, v1))
self.assertAllEqual(deq_elem, sorted(elem))
for e, dv0, dv1 in zip(deq_elem, deq_value_0, deq_value_1):
self.assertTrue((dv0, dv1) in allowed[e])
missed.remove((dv0, dv1))
self.assertEqual(missed, set())
def testRoundTripFillsCapacityMultiThreadedEnqueueAndDequeue(self):
# We need each thread to keep its own device stack or the device scopes
# won't be properly nested.
ops.get_default_graph().switch_to_thread_local()
with self.cached_session() as sess:
q = data_flow_ops.PriorityQueue(10, (dtypes.int64), (()))
num_threads = 40
enqueue_counts = np.random.randint(10, size=num_threads)
enqueue_values = [
np.random.randint(
5, size=count) for count in enqueue_counts
]
enqueue_ops = [
q.enqueue_many((values, values)) for values in enqueue_values
]
shuffled_counts = copy.deepcopy(enqueue_counts)
random.shuffle(shuffled_counts)
dequeue_ops = [q.dequeue_many(count) for count in shuffled_counts]
all_enqueued_values = np.hstack(enqueue_values)
# Run one producer thread for each element in elems.
def enqueue(enqueue_op):
self.evaluate(enqueue_op)
dequeued = []
def dequeue(dequeue_op):
(dequeue_indices, dequeue_values) = self.evaluate(dequeue_op)
self.assertAllEqual(dequeue_indices, dequeue_values)
dequeued.extend(dequeue_indices)
enqueue_threads = [
self.checkedThread(
target=enqueue, args=(op,)) for op in enqueue_ops
]
dequeue_threads = [
self.checkedThread(
target=dequeue, args=(op,)) for op in dequeue_ops
]
# Dequeue and check
for t in dequeue_threads:
t.start()
for t in enqueue_threads:
t.start()
for t in enqueue_threads:
t.join()
for t in dequeue_threads:
t.join()
self.assertAllEqual(sorted(dequeued), sorted(all_enqueued_values))
def testRoundTripInsertManyMultiThreadedReadManyMultithreadedSorts(self):
# We need each thread to keep its own device stack or the device scopes
# won't be properly nested.
ops.get_default_graph().switch_to_thread_local()
with self.cached_session() as sess:
q = data_flow_ops.PriorityQueue(2000, (dtypes.int64), (()))
num_threads = 40
enqueue_counts = np.random.randint(10, size=num_threads)
enqueue_values = [
np.random.randint(
5, size=count) for count in enqueue_counts
]
enqueue_ops = [
q.enqueue_many((values, values)) for values in enqueue_values
]
shuffled_counts = copy.deepcopy(enqueue_counts)
random.shuffle(shuffled_counts)
dequeue_ops = [q.dequeue_many(count) for count in shuffled_counts]
all_enqueued_values = np.hstack(enqueue_values)
dequeue_wait = threading.Condition()
# Run one producer thread for each element in elems.
def enqueue(enqueue_op):
self.evaluate(enqueue_op)
def dequeue(dequeue_op, dequeued):
(dequeue_indices, dequeue_values) = self.evaluate(dequeue_op)
self.assertAllEqual(dequeue_indices, dequeue_values)
dequeue_wait.acquire()
dequeued.extend(dequeue_indices)
dequeue_wait.release()
dequeued = []
enqueue_threads = [
self.checkedThread(
target=enqueue, args=(op,)) for op in enqueue_ops
]
dequeue_threads = [
self.checkedThread(
target=dequeue, args=(op, dequeued)) for op in dequeue_ops
]
for t in enqueue_threads:
t.start()
for t in enqueue_threads:
t.join()
# Dequeue and check
for t in dequeue_threads:
t.start()
for t in dequeue_threads:
t.join()
# We can't guarantee full sorting because we can't guarantee
# that the dequeued.extend() call runs immediately after the
# self.evaluate() call. Here we're just happy everything came out.
self.assertAllEqual(set(dequeued), set(all_enqueued_values))
def testRoundTripInsertManyMultiThreadedReadOnceSorts(self):
# We need each thread to keep its own device stack or the device scopes
# won't be properly nested.
ops.get_default_graph().switch_to_thread_local()
with self.cached_session() as sess:
q = data_flow_ops.PriorityQueue(2000, (dtypes.string, dtypes.string), (
(), ()))
elem = np.random.randint(-5, 5, size=100).astype(np.int64)
side_value_0 = np.random.rand(100).astype(bytes)
side_value_1 = np.random.rand(100).astype(bytes)
batch = 5
enqueue_ops = [
q.enqueue_many((elem[i * batch:(i + 1) * batch],
side_value_0[i * batch:(i + 1) * batch],
side_value_1[i * batch:(i + 1) * batch]))
for i in range(20)
]
# Run one producer thread for each element in elems.
def enqueue(enqueue_op):
self.evaluate(enqueue_op)
dequeue_op = q.dequeue_many(100)
enqueue_threads = [
self.checkedThread(
target=enqueue, args=(op,)) for op in enqueue_ops
]
for t in enqueue_threads:
t.start()
deq_elem, deq_value_0, deq_value_1 = self.evaluate(dequeue_op)
for t in enqueue_threads:
t.join()
allowed = {}
missed = set()
for e, v0, v1 in zip(elem, side_value_0, side_value_1):
if e not in allowed:
allowed[e] = set()
allowed[e].add((v0, v1))
missed.add((v0, v1))
self.assertAllEqual(deq_elem, sorted(elem))
for e, dv0, dv1 in zip(deq_elem, deq_value_0, deq_value_1):
self.assertTrue((dv0, dv1) in allowed[e])
missed.remove((dv0, dv1))
self.assertEqual(missed, set())
def testRoundTripInsertOnceReadOnceSorts(self):
with self.cached_session() as sess:
q = data_flow_ops.PriorityQueue(2000, (dtypes.string, dtypes.string), (
(), ()))
elem = np.random.randint(-100, 100, size=1000).astype(np.int64)
side_value_0 = np.random.rand(1000).astype(bytes)
side_value_1 = np.random.rand(1000).astype(bytes)
q.enqueue_many((elem, side_value_0, side_value_1)).run()
deq = q.dequeue_many(1000)
deq_elem, deq_value_0, deq_value_1 = self.evaluate(deq)
allowed = {}
for e, v0, v1 in zip(elem, side_value_0, side_value_1):
if e not in allowed:
allowed[e] = set()
allowed[e].add((v0, v1))
self.assertAllEqual(deq_elem, sorted(elem))
for e, dv0, dv1 in zip(deq_elem, deq_value_0, deq_value_1):
self.assertTrue((dv0, dv1) in allowed[e])
def testRoundTripInsertOnceReadManySorts(self):
with self.cached_session():
q = data_flow_ops.PriorityQueue(2000, (dtypes.int64), (()))
elem = np.random.randint(-100, 100, size=1000).astype(np.int64)
q.enqueue_many((elem, elem)).run()
deq_values = np.hstack([q.dequeue_many(100)[0].eval() for _ in range(10)])
self.assertAllEqual(deq_values, sorted(elem))
def testRoundTripInsertOnceReadOnceLotsSorts(self):
with self.cached_session():
q = data_flow_ops.PriorityQueue(2000, (dtypes.int64), (()))
elem = np.random.randint(-100, 100, size=1000).astype(np.int64)
q.enqueue_many((elem, elem)).run()
dequeue_op = q.dequeue()
deq_values = np.hstack([dequeue_op[0].eval() for _ in range(1000)])
self.assertAllEqual(deq_values, sorted(elem))
def testInsertingNonInt64Fails(self):
with self.cached_session():
q = data_flow_ops.PriorityQueue(2000, (dtypes.string), (()))
with self.assertRaises(TypeError):
q.enqueue_many((["a", "b", "c"], ["a", "b", "c"])).run()
def testInsertingNonScalarFails(self):
with self.cached_session() as sess:
input_priority = array_ops.placeholder(dtypes.int64)
input_other = array_ops.placeholder(dtypes.string)
q = data_flow_ops.PriorityQueue(2000, (dtypes.string,), (()))
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
r"Shape mismatch in tuple component 0. Expected \[\], got \[2\]"):
sess.run([q.enqueue((input_priority, input_other))],
feed_dict={
input_priority: np.array(
[0, 2], dtype=np.int64),
input_other: np.random.rand(3, 5).astype(bytes)
})
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
r"Shape mismatch in tuple component 0. Expected \[2\], got \[2,2\]"):
sess.run(
[q.enqueue_many((input_priority, input_other))],
feed_dict={
input_priority: np.array(
[[0, 2], [3, 4]], dtype=np.int64),
input_other: np.random.rand(2, 3).astype(bytes)
})
if __name__ == "__main__":
test.main()
tensorflow/python/kernel_tests/data_structures/stack_ops_test.py
+299
View File
@@ -0,0 +1,299 @@
# Copyright 2015 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 tensorflow.ops.stack_ops."""
import numpy as np
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.ops import gen_data_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import while_loop
from tensorflow.python.platform import test
class StackOpTest(test.TestCase):
def _testStackPushPop(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
h = gen_data_flow_ops.stack_v2(
-1, elem_type=dtypes.float32, stack_name="foo")
c = gen_data_flow_ops.stack_push_v2(h, [[4.0, 5.0]])
with ops.control_dependencies([c]):
c1 = gen_data_flow_ops.stack_pop_v2(h, dtypes.float32)
self.assertAllClose([[4.0, 5.0]], self.evaluate(c1))
@test_util.run_deprecated_v1
def testStackPushPop(self):
self._testStackPushPop(use_gpu=False)
self._testStackPushPop(use_gpu=True)
def _testStackPushPopSwap(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
a = np.arange(2000)
x = constant_op.constant(a, dtype=dtypes.float32)
h = gen_data_flow_ops.stack_v2(
-1, elem_type=dtypes.float32, stack_name="foo")
c = gen_data_flow_ops.stack_push_v2(h, x, swap_memory=True)
with ops.control_dependencies([c]):
c1 = gen_data_flow_ops.stack_pop_v2(h, dtypes.float32)
self.assertAllClose(a, self.evaluate(c1))
@test_util.run_deprecated_v1
def testStackPushPopSwap(self):
self._testStackPushPopSwap(use_gpu=False)
self._testStackPushPopSwap(use_gpu=True)
def _testStackWhileSwap(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
n = constant_op.constant(0)
h = gen_data_flow_ops.stack_v2(
-1, elem_type=dtypes.float32, stack_name="foo")
def c(x):
return math_ops.less(x, 10)
def b(x):
with ops.control_dependencies([x]):
a = constant_op.constant(np.ones(2000), dtype=dtypes.float32)
v = gen_data_flow_ops.stack_push_v2(h, a, swap_memory=True)
with ops.control_dependencies([v]):
return math_ops.add(x, 1)
r = while_loop.while_loop(c, b, [n])
v = constant_op.constant(np.zeros(2000), dtype=dtypes.float32)
def c1(x, y):
del y
return math_ops.greater(x, 0)
def b1(x, y):
nx = math_ops.subtract(x, 1)
ny = y + gen_data_flow_ops.stack_pop_v2(h, dtypes.float32)
return [nx, ny]
_, ry = while_loop.while_loop(
c1, b1, [r, v],
[r.get_shape(), tensor_shape.unknown_shape()])
self.assertAllClose(np.ones(2000) * 10.0, self.evaluate(ry))
@test_util.run_v1_only("b/120545219")
def testStackWhileSwap(self):
self._testStackWhileSwap(use_gpu=False)
self._testStackWhileSwap(use_gpu=True)
def _testMultiStack(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
h1 = gen_data_flow_ops.stack_v2(
-1, elem_type=dtypes.float32, stack_name="foo")
c1 = gen_data_flow_ops.stack_push_v2(h1, 4.0)
with ops.control_dependencies([c1]):
c1 = gen_data_flow_ops.stack_pop_v2(h1, dtypes.float32)
h2 = gen_data_flow_ops.stack_v2(
-1, elem_type=dtypes.float32, stack_name="bar")
c2 = gen_data_flow_ops.stack_push_v2(h2, 5.0)
with ops.control_dependencies([c2]):
c2 = gen_data_flow_ops.stack_pop_v2(h2, dtypes.float32)
r = c1 + c2
self.assertAllClose(9.0, self.evaluate(r))
@test_util.run_deprecated_v1
def testMultiStack(self):
self._testMultiStack(use_gpu=False)
self._testMultiStack(use_gpu=True)
def _testSameNameStacks(self, use_gpu):
"""Different stacks with the same name do not interfere."""
with self.cached_session(use_gpu=use_gpu) as sess:
h1 = gen_data_flow_ops.stack_v2(
-1, elem_type=dtypes.float32, stack_name="foo")
h2 = gen_data_flow_ops.stack_v2(
-1, elem_type=dtypes.float32, stack_name="foo")
c1 = gen_data_flow_ops.stack_push_v2(h1, 4.0)
with ops.control_dependencies([c1]):
c2 = gen_data_flow_ops.stack_push_v2(h2, 5.0)
with ops.control_dependencies([c2]):
pop1 = gen_data_flow_ops.stack_pop_v2(h1, dtypes.float32)
pop2 = gen_data_flow_ops.stack_pop_v2(h2, dtypes.float32)
out1, out2 = self.evaluate([pop1, pop2])
self.assertAllClose(out1, 4.0)
self.assertAllClose(out2, 5.0)
@test_util.run_deprecated_v1
def testSameNameStacks(self):
self._testSameNameStacks(use_gpu=False)
self._testSameNameStacks(use_gpu=True)
def _testCloseStack(self, use_gpu):
with self.cached_session(use_gpu=use_gpu) as sess:
h = gen_data_flow_ops.stack_v2(
-1, elem_type=dtypes.float32, stack_name="foo")
c1 = gen_data_flow_ops.stack_close_v2(h)
self.evaluate(c1)
@test_util.run_deprecated_v1
def testCloseStack(self):
self._testCloseStack(use_gpu=False)
self._testCloseStack(use_gpu=True)
def _testPushCloseStack(self, use_gpu):
with self.cached_session(use_gpu=use_gpu) as sess:
h = gen_data_flow_ops.stack_v2(
-1, elem_type=dtypes.float32, stack_name="foo")
c = gen_data_flow_ops.stack_push_v2(h, [[4.0, 5.0]])
with ops.control_dependencies([c]):
c1 = gen_data_flow_ops.stack_close_v2(h)
self.evaluate(c1)
@test_util.run_deprecated_v1
def testPushCloseStack(self):
self._testPushCloseStack(use_gpu=False)
self._testPushCloseStack(use_gpu=True)
class StackOpRefTest(test.TestCase):
"""Tests for deprecated non-resource variant of stack ops."""
def _testStackPushPop(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
h = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo")
c = gen_data_flow_ops.stack_push(h, [[4.0, 5.0]])
with ops.control_dependencies([c]):
c1 = gen_data_flow_ops.stack_pop(h, dtypes.float32)
self.assertAllClose([[4.0, 5.0]], self.evaluate(c1))
@test_util.run_deprecated_v1
def testStackPushPop(self):
self._testStackPushPop(use_gpu=False)
self._testStackPushPop(use_gpu=True)
def _testStackPushPopSwap(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
a = np.arange(2000)
x = constant_op.constant(a, dtype=dtypes.float32)
h = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo")
c = gen_data_flow_ops.stack_push(h, x, swap_memory=True)
with ops.control_dependencies([c]):
c1 = gen_data_flow_ops.stack_pop(h, dtypes.float32)
self.assertAllClose(a, self.evaluate(c1))
@test_util.run_deprecated_v1
def testStackPushPopSwap(self):
self._testStackPushPopSwap(use_gpu=False)
self._testStackPushPopSwap(use_gpu=True)
def _testMultiStack(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
h1 = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo")
c1 = gen_data_flow_ops.stack_push(h1, 4.0)
with ops.control_dependencies([c1]):
c1 = gen_data_flow_ops.stack_pop(h1, dtypes.float32)
h2 = gen_data_flow_ops._stack(dtypes.float32, stack_name="bar")
c2 = gen_data_flow_ops.stack_push(h2, 5.0)
with ops.control_dependencies([c2]):
c2 = gen_data_flow_ops.stack_pop(h2, dtypes.float32)
r = c1 + c2
self.assertAllClose(9.0, self.evaluate(r))
def _testStackWhileSwap(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
n = constant_op.constant(0)
h = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo")
def c(x):
return math_ops.less(x, 10)
def b(x):
with ops.control_dependencies([x]):
a = constant_op.constant(np.ones(2000), dtype=dtypes.float32)
v = gen_data_flow_ops.stack_push(h, a, swap_memory=True)
with ops.control_dependencies([v]):
return math_ops.add(x, 1)
r = while_loop.while_loop(c, b, [n])
v = constant_op.constant(np.zeros(2000), dtype=dtypes.float32)
def c1(x, y):
del y
return math_ops.greater(x, 0)
def b1(x, y):
nx = math_ops.subtract(x, 1)
ny = y + gen_data_flow_ops.stack_pop(h, dtypes.float32)
return [nx, ny]
_, ry = while_loop.while_loop(
c1, b1, [r, v],
[r.get_shape(), tensor_shape.unknown_shape()])
self.assertAllClose(np.ones(2000) * 10.0, self.evaluate(ry))
@test_util.run_v1_only("b/120545219")
def testStackWhileSwap(self):
self._testStackWhileSwap(use_gpu=False)
self._testStackWhileSwap(use_gpu=True)
@test_util.run_deprecated_v1
def testMultiStack(self):
self._testMultiStack(use_gpu=False)
self._testMultiStack(use_gpu=True)
def _testSameNameStacks(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
h1 = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo")
c1 = gen_data_flow_ops.stack_push(h1, 4.0)
h2 = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo")
c2 = gen_data_flow_ops.stack_push(h2, 5.0)
_ = c1 + c2
self.assertNotEqual(self.evaluate(h1)[1], self.evaluate(h2)[1])
@test_util.run_deprecated_v1
def testSameNameStacks(self):
self._testSameNameStacks(use_gpu=False)
self._testSameNameStacks(use_gpu=True)
def _testCloseStack(self, use_gpu):
with self.cached_session(use_gpu=use_gpu) as sess:
h = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo")
c1 = gen_data_flow_ops.stack_close(h)
self.evaluate(c1)
@test_util.run_deprecated_v1
def testCloseStack(self):
self._testCloseStack(use_gpu=False)
self._testCloseStack(use_gpu=True)
def _testPushCloseStack(self, use_gpu):
with self.cached_session(use_gpu=use_gpu) as sess:
h = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo")
c = gen_data_flow_ops.stack_push(h, [[4.0, 5.0]])
with ops.control_dependencies([c]):
c1 = gen_data_flow_ops.stack_close(h)
self.evaluate(c1)
@test_util.run_deprecated_v1
def testPushCloseStack(self):
self._testPushCloseStack(use_gpu=False)
self._testPushCloseStack(use_gpu=True)
if __name__ == "__main__":
test.main()
tensorflow/python/kernel_tests/data_structures/stage_op_test.py
+299
View File
@@ -0,0 +1,299 @@
# Copyright 2017 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.
# ==============================================================================
import queue
import threading
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
TIMEOUT = 1
class StageTest(test.TestCase):
@test_util.run_deprecated_v1
def testSimple(self):
with ops.Graph().as_default() as G:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.StagingArea([dtypes.float32])
stage = stager.put([v])
y = stager.get()
y = math_ops.reduce_max(math_ops.matmul(y, y))
G.finalize()
with self.session(graph=G) as sess:
sess.run(stage, feed_dict={x: -1})
for i in range(10):
_, yval = sess.run([stage, y], feed_dict={x: i})
self.assertAllClose(4 * (i - 1) * (i - 1) * 128, yval, rtol=1e-4)
@test_util.run_deprecated_v1
def testMultiple(self):
with ops.Graph().as_default() as G:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.StagingArea([dtypes.float32, dtypes.float32])
stage = stager.put([x, v])
z, y = stager.get()
y = math_ops.reduce_max(z * math_ops.matmul(y, y))
G.finalize()
with self.session(graph=G) as sess:
sess.run(stage, feed_dict={x: -1})
for i in range(10):
_, yval = sess.run([stage, y], feed_dict={x: i})
self.assertAllClose(
4 * (i - 1) * (i - 1) * (i - 1) * 128, yval, rtol=1e-4)
@test_util.run_deprecated_v1
def testDictionary(self):
with ops.Graph().as_default() as G:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.StagingArea(
[dtypes.float32, dtypes.float32],
shapes=[[], [128, 128]],
names=['x', 'v'])
stage = stager.put({'x': x, 'v': v})
ret = stager.get()
z = ret['x']
y = ret['v']
y = math_ops.reduce_max(z * math_ops.matmul(y, y))
G.finalize()
with self.session(graph=G) as sess:
sess.run(stage, feed_dict={x: -1})
for i in range(10):
_, yval = sess.run([stage, y], feed_dict={x: i})
self.assertAllClose(
4 * (i - 1) * (i - 1) * (i - 1) * 128, yval, rtol=1e-4)
def testColocation(self):
gpu_dev = test.gpu_device_name()
with ops.Graph().as_default() as G:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32)
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(gpu_dev):
stager = data_flow_ops.StagingArea([dtypes.float32])
y = stager.put([v])
expected_name = gpu_dev if 'gpu' not in gpu_dev else '/device:GPU:0'
self.assertEqual(y.device, expected_name)
with ops.device('/cpu:0'):
x = stager.get()[0]
self.assertEqual(x.device, '/device:CPU:0')
G.finalize()
@test_util.run_deprecated_v1
def testPeek(self):
with ops.Graph().as_default() as G:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.int32, name='x')
p = array_ops.placeholder(dtypes.int32, name='p')
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.StagingArea(
[
dtypes.int32,
], shapes=[[]])
stage = stager.put([x])
peek = stager.peek(p)
ret = stager.get()
G.finalize()
with self.session(graph=G) as sess:
for i in range(10):
sess.run(stage, feed_dict={x: i})
for i in range(10):
self.assertTrue(sess.run(peek, feed_dict={p: i}) == [i])
def testPeekBadIndex(self):
stager = data_flow_ops.StagingArea([
dtypes.int32,
], shapes=[[10]])
stager.put([array_ops.zeros([10], dtype=dtypes.int32)])
with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
'must be scalar'):
self.evaluate(stager.peek([]))
@test_util.run_deprecated_v1
def testSizeAndClear(self):
with ops.Graph().as_default() as G:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.float32, name='x')
v = 2. * (array_ops.zeros([128, 128]) + x)
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.StagingArea(
[dtypes.float32, dtypes.float32],
shapes=[[], [128, 128]],
names=['x', 'v'])
stage = stager.put({'x': x, 'v': v})
ret = stager.get()
size = stager.size()
clear = stager.clear()
G.finalize()
with self.session(graph=G) as sess:
sess.run(stage, feed_dict={x: -1})
self.assertEqual(sess.run(size), 1)
sess.run(stage, feed_dict={x: -1})
self.assertEqual(sess.run(size), 2)
sess.run(clear)
self.assertEqual(sess.run(size), 0)
@test_util.run_deprecated_v1
def testCapacity(self):
self.skipTest('b/123423516 this test is flaky on gpu.')
capacity = 3
with ops.Graph().as_default() as G:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.int32, name='x')
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.StagingArea(
[
dtypes.int32,
], capacity=capacity, shapes=[[]])
stage = stager.put([x])
ret = stager.get()
size = stager.size()
G.finalize()
value_queue = queue.Queue()
n = 8
with self.session(graph=G) as sess:
# Stage data in a separate thread which will block when it hits the
# staging area's capacity and thus not fill the value_queue with n tokens
def thread_run():
for i in range(n):
sess.run(stage, feed_dict={x: i})
value_queue.put(0)
t = threading.Thread(target=thread_run)
t.daemon = True
t.start()
# Get tokens from the value_queue until a timeout occurs
try:
for i in range(n):
value_queue.get(timeout=TIMEOUT)
except queue.Empty:
pass
# Should've timed out on the iteration 'capacity'
if not i == capacity:
self.fail("Expected to timeout on iteration '{}' "
"but instead timed out on iteration '{}' "
"Staging Area size is '{}' and configured "
"capacity is '{}'.".format(capacity, i, sess.run(size),
capacity))
# Should have capacity elements in the staging area
self.assertTrue(sess.run(size) == capacity)
# Clear the staging area completely
for i in range(n):
self.assertTrue(sess.run(ret) == [i])
# It should now be empty
self.assertTrue(sess.run(size) == 0)
@test_util.run_deprecated_v1
def testMemoryLimit(self):
memory_limit = 512 * 1024 # 512K
chunk = 200 * 1024 # 256K
capacity = memory_limit // chunk
with ops.Graph().as_default() as G:
with ops.device('/cpu:0'):
x = array_ops.placeholder(dtypes.uint8, name='x')
with ops.device(test.gpu_device_name()):
stager = data_flow_ops.StagingArea(
[
dtypes.uint8,
], memory_limit=memory_limit, shapes=[[]])
stage = stager.put([x])
ret = stager.get()
size = stager.size()
G.finalize()
value_queue = queue.Queue()
n = 8
with self.session(graph=G) as sess:
# Stage data in a separate thread which will block when it hits the
# staging area's capacity and thus not fill the value_queue with n tokens
def thread_run():
for i in range(n):
sess.run(stage, feed_dict={x: np.full(chunk, i, dtype=np.uint8)})
value_queue.put(0)
t = threading.Thread(target=thread_run)
t.daemon = True
t.start()
# Get tokens from the value_queue until a timeout occurs
try:
for i in range(n):
value_queue.get(timeout=TIMEOUT)
except queue.Empty:
pass
# Should've timed out on the iteration 'capacity'
if not i == capacity:
self.fail("Expected to timeout on iteration '{}' "
"but instead timed out on iteration '{}' "
"Staging Area size is '{}' and configured "
"capacity is '{}'.".format(capacity, i, sess.run(size),
capacity))
# Should have capacity elements in the staging area
self.assertTrue(sess.run(size) == capacity)
# Clear the staging area completely
for i in range(n):
self.assertTrue(np.all(sess.run(ret)[0] == i))
self.assertTrue(sess.run(size) == 0)
if __name__ == '__main__':
test.main()
tensorflow/python/kernel_tests/data_structures/tensor_array_ops_test.py
+1908
View File
File diff suppressed because it is too large Load Diff