def dataset_producer(t):
ragged_ds = dataset_ops.Dataset.from_tensor_slices(t).batch(
2, drop_remainder)
it = multi_device_iterator_ops.MultiDeviceIterator(
ragged_ds, ['GPU:0'])
with ops.device_v2('GPU:0'):
return it.get_next_as_optional()
def testDefaultDeviceInsideFunctionWithScope(self, distribution,
run_functions_eagerly):
def_function.run_functions_eagerly(run_functions_eagerly)
expected_device = (device_util.canonicalize("cpu:0")
if run_functions_eagerly else "")
with distribution.scope():
with ops.device_v2("cpu:0"):
@def_function.function
def add():
one = array_ops.ones([])
self.assertEqual(expected_device, one.device)
return one + 1
add()
def testDefaultDeviceInsideFunctionWithScope(
self, distribution, run_functions_eagerly):
def_function.run_functions_eagerly(run_functions_eagerly)
try:
worker = distribution.extended.worker_devices[0]
except RuntimeError:
worker = None
expected_device = (device_util.canonicalize("cpu:0", worker)
if run_functions_eagerly else "")
with distribution.scope():
with ops.device_v2("cpu:0"):
@def_function.function
def add():
one = array_ops.ones([])
self.assertEqual(expected_device, one.device)
return one + 1
add()
def bm_layer_implementation(self, batch_size):
with ops.device_v2("/gpu:0"):
img = keras.Input(shape=(256, 256, 3), dtype=dtypes.float32)
preprocessor = keras.Sequential([
image_preprocessing.Resizing(224, 224),
image_preprocessing.RandomCrop(height=224, width=224),
image_preprocessing.RandomRotation(factor=(.2, .4)),
image_preprocessing.RandomFlip(mode="horizontal"),
image_preprocessing.RandomZoom(.2, .2)
])
_ = preprocessor(img)
num_repeats = 5
starts = []
ends = []
for _ in range(num_repeats):
ds = dataset_ops.Dataset.from_tensor_slices(
np.random.random((batch_size, 256, 256, 3)))
ds = ds.shuffle(batch_size * 100)
ds = ds.batch(batch_size)
ds = ds.prefetch(batch_size)
starts.append(time.time())
count = 0
# Benchmarked code begins here.
for i in ds:
_ = preprocessor(i)
count += 1
# Benchmarked code ends here.
ends.append(time.time())
avg_time = np.mean(np.array(ends) - np.array(starts)) / count
name = "image_preprocessing|batch_%s" % batch_size
baseline = self.run_dataset_implementation(batch_size)
extras = {
"dataset implementation baseline": baseline,
"delta seconds": (baseline - avg_time),
"delta percent": ((baseline - avg_time) / baseline) * 100
}
self.report_benchmark(iters=num_repeats,
wall_time=avg_time,
extras=extras,
name=name)
def _embedding_lookup_and_transform(params,
ids,
partition_strategy="mod",
name=None,
max_norm=None,
transform_fn=None):
"""Helper function for embedding_lookup and _compute_sampled_logits.
This function is a generalization of embedding_lookup that optionally
applies a caller-specified transformation to each embedding. This is
done through the `transform_fn` argument. If provided, the function is
applied to each partitioned tensor of retrieved embeddings, colocated
with the embeddings. This function will be called with a single `Tensor`
argument of the same type as the `params` tensor and should return a
`Tensor`. The shape of the argument will be the same as `params` except
for the size of the first dimension. The first dimension of the result's
shape must be the same size as the argument's.
Args:
params: See embedding_lookup.
ids: See embedding_lookup.
partition_strategy: See embedding_lookup.
name: See embedding_lookup.
max_norm: See embedding_lookup.
transform_fn: An optional function to apply to each retrieved embedding. If
max_norm is provided, transform_fn is applied to the norm-limited
embeddings.
Returns:
See embedding_lookup for details.
Raises:
ValueError: If `params` is empty.
"""
if params is None:
raise ValueError("params must be specified")
if isinstance(params, (list, tuple)) and not params:
raise ValueError("Need at least one param")
if isinstance(params, variables.PartitionedVariable):
params = list(params) # Iterate to get the underlying Variables.
if not isinstance(params, list):
params = [params]
with ops.name_scope(name, "embedding_lookup", params + [ids]) as name:
np = len(params) # Number of partitions
# Preserve the resource variable status to avoid accidental dense reads.
if not any(
isinstance(p, resource_variable_ops.BaseResourceVariable)
for p in params):
params = ops.convert_n_to_tensor_or_indexed_slices(params, name="params")
ids = ops.convert_to_tensor(ids, name="ids")
if np == 1 and (not transform_fn or ids.get_shape().ndims == 1):
with ops.colocate_with(params[0]):
result = _clip(
array_ops.gather(params[0], ids, name=name), ids, max_norm)
if transform_fn:
result = transform_fn(result)
# Make sure the final result does not have colocation constraints on the
# params. Similar to the case np > 1 where parallel_dynamic_stitch is
# outside the scioe of all with ops.colocate_with(params[p]).
return array_ops.identity(result)
else:
# Flatten the ids. There are two cases where we need to do this.
# - There is more than one params tensor.
# - There is a transform_fn and ids is not statically known to be 1-D.
# We must flatten in this case because transform_fn expects a flat
# tensor of embeddings.
flat_ids = array_ops.reshape(ids, [-1])
original_indices = math_ops.range(array_ops.size(flat_ids))
# Create p_assignments and set new_ids depending on the strategy.
if partition_strategy == "mod":
p_assignments = flat_ids % np
new_ids = flat_ids // np
elif partition_strategy == "div":
# Compute num_total_ids as the sum of dim-0 of params, then assign to
# partitions based on a constant number of ids per partition. Optimize
# if we already know the full shape statically.
dim_0_size = tensor_shape.Dimension(
tensor_shape.dimension_value(params[0].get_shape()[0]))
for p in xrange(1, np):
dim_0_size += tensor_shape.Dimension(
tensor_shape.dimension_value(params[p].get_shape()[0]))
if dim_0_size.value:
num_total_ids = constant_op.constant(dim_0_size.value, flat_ids.dtype)
else:
dim_0_sizes = []
for p in xrange(np):
param_p_dim = tensor_shape.dimension_value(params[p].get_shape()[0])
if param_p_dim is not None:
dim_0_sizes.append(param_p_dim)
else:
with ops.colocate_with(params[p]):
dim_0_sizes.append(array_ops.shape(params[p])[0])
num_total_ids = math_ops.reduce_sum(
math_ops.cast(array_ops.stack(dim_0_sizes), flat_ids.dtype))
ids_per_partition = num_total_ids // np
extras = num_total_ids % np
p_assignments = math_ops.maximum(flat_ids // (ids_per_partition + 1),
(flat_ids - extras) //
ids_per_partition)
# Emulate a conditional using a boolean indicator tensor
new_ids = array_ops.where(p_assignments < extras,
flat_ids % (ids_per_partition + 1),
(flat_ids - extras) % ids_per_partition)
else:
raise ValueError("Unrecognized partition strategy: " +
partition_strategy)
# Cast partition assignments to int32 for use in dynamic_partition.
# There really should not be more than 2^32 partitions.
p_assignments = math_ops.cast(p_assignments, dtypes.int32)
# Partition list of ids based on assignments into np separate lists
gather_ids = data_flow_ops.dynamic_partition(new_ids, p_assignments, np)
# Similarly, partition the original indices.
pindices = data_flow_ops.dynamic_partition(original_indices,
p_assignments, np)
# Do np separate lookups, finding embeddings for plist[p] in params[p]
partitioned_result = []
for p in xrange(np):
pids = gather_ids[p]
with ops.device_v2(None):
with ops.colocate_with(params[p]):
result = array_ops.gather(params[p], pids)
if transform_fn:
# If transform_fn is provided, the clip_by_norm precedes
# the transform and hence must be co-located. See below
# for the counterpart if transform_fn is not provided.
result = transform_fn(_clip(result, pids, max_norm))
partitioned_result.append(result)
# Stitch these back together
ret = data_flow_ops.parallel_dynamic_stitch(
pindices, partitioned_result, name=name)
# Determine the static element shape.
if transform_fn is None:
element_shape_s = params[0].get_shape()[1:]
for p in params[1:]:
element_shape_s = element_shape_s.merge_with(p.get_shape()[1:])
else:
element_shape_s = ret.get_shape()[1:]
# Compute the dynamic element shape.
if element_shape_s.is_fully_defined():
element_shape_d = element_shape_s
elif transform_fn is None:
# It's important that we compute params[0].shape on the right device
# to avoid data motion.
with ops.colocate_with(params[0]):
params_shape = array_ops.shape(params[0])
element_shape_d = params_shape[1:]
else:
element_shape_d = array_ops.shape(ret)[1:]
# Reshape to reverse the flattening of ids.
ret = array_ops.reshape(
ret, array_ops.concat([array_ops.shape(ids), element_shape_d], 0))
# Normally the reshape is sufficient, but setting shape explicitly
# teaches shape inference that params[1:].get_shape() matters
# (in the case that transform_fn is None).
ret.set_shape(ids.get_shape().concatenate(element_shape_s))
if not transform_fn:
# If transform_fn was provided, the clip_by_norm was done above.
ret = _clip(ret, ids, max_norm)
return ret
def testBadIndicesCPU(self):
with ops.device_v2("cpu:0"):
params = [[0, 1, 2], [3, 4, 5]]
with self.assertRaisesOpError(
r"indices\[0\] = 7 is not in \[0, 2\)"):
self.evaluate(array_ops.batch_gather(params, [7]))
