def testUsingInfeedQueueWithRegularizer(self):
"""Test that Layer regularizers can reference data created in loops."""
def make_regularizer(scale):
return lambda inputs: scale * math_ops.reduce_sum(
math_ops.square(inputs))
def training_step(inputs, scale):
outputs = convolutional.conv2d(
inputs,
filters=16,
kernel_size=(3, 3),
data_format="channels_first",
kernel_regularizer=make_regularizer(scale))
loss = math_ops.reduce_mean(math_ops.square(outputs))
return loss.op
inputs = array_ops.zeros(shape=(128, 32, 32, 16))
scale = array_ops.ones(shape=())
infeed = tpu_feed.InfeedQueue(
tuple_types=[dtypes.float32, dtypes.float32],
tuple_shapes=[inputs.shape, scale.shape])
def loop():
return training_loop.repeat(5, training_step, infeed_queue=infeed)
# This should not throw an error.
tpu.rewrite(loop)
def testUsingInfeedQueueWithRegularizer(self):
"""Test that Layer regularizers can reference data created in loops."""
def make_regularizer(scale):
return lambda inputs: scale * math_ops.reduce_sum(math_ops.square(inputs))
def training_step(inputs, scale):
outputs = convolutional.conv2d(
inputs,
filters=16,
kernel_size=(3, 3),
data_format="channels_first",
kernel_regularizer=make_regularizer(scale))
loss = math_ops.reduce_mean(math_ops.square(outputs))
return loss.op
inputs = array_ops.zeros(shape=(128, 32, 32, 16))
scale = array_ops.ones(shape=())
infeed = tpu_feed.InfeedQueue(
tuple_types=[dtypes.float32, dtypes.float32],
tuple_shapes=[inputs.shape, scale.shape])
def loop():
return training_loop.repeat(5, training_step, infeed_queue=infeed)
# This should not throw an error.
tpu.rewrite(loop)
def run_on_device(self, model_fn, model_inputs, device):
"""Runs `model_fn` on the given device.
Raises an exception if no such device is available. `model_fn` should
return one or more tensors as a list or tuple.
Args:
model_fn: Function returning one or more tensors.
model_inputs: An iterable of Numpy arrays or scalars.
These will be passed as arguments to `model_fn`.
device: Device to run on. One of ("tpu", "gpu", "cpu").
Returns:
Output from the model function.
"""
def _make_placeholders():
return dict([(gen_array_ops.placeholder_with_default(v, v.shape), v)
for v in model_inputs])
if device == "tpu":
with self.test_session(graph=ops.Graph()) as sess:
placeholders = _make_placeholders()
tpu_computation = tpu.rewrite(model_fn, placeholders.keys())
sess.run(tpu.initialize_system())
sess.run(variables.global_variables_initializer())
result = sess.run(tpu_computation, placeholders)
sess.run(tpu.shutdown_system())
# TODO(b/36891278): supports non-flat returns lists in tpu.rewrite().
if len(result) == 1:
return result[0]
return result
elif device == "gpu":
with self.test_session(graph=ops.Graph(), use_gpu=True) as sess:
placeholders = _make_placeholders()
sess.run(variables.global_variables_initializer())
return sess.run(model_fn(placeholders.keys()), placeholders)
elif device == "cpu":
# TODO(power) -- will this interact poorly with cached GPU sessions?
with self.test_session(graph=ops.Graph(), use_gpu=False) as sess:
placeholders = _make_placeholders()
sess.run(variables.global_variables_initializer())
return sess.run(model_fn(placeholders.keys()), placeholders)
def run_on_device(self, model_fn, model_inputs, device):
"""Runs `model_fn` on the given device.
Raises an exception if no such device is available. `model_fn` should
return one or more tensors as a list or tuple.
Args:
model_fn: Function returning one or more tensors.
model_inputs: An iterable of Numpy arrays or scalars.
These will be passed as arguments to `model_fn`.
device: Device to run on. One of ("tpu", "gpu", "cpu").
Returns:
Output from the model function.
"""
def _make_placeholders():
return dict(
[(gen_array_ops.placeholder_with_default(v, v.shape), v)
for v in model_inputs])
if device == "tpu":
with self.test_session(graph=ops.Graph()) as sess:
placeholders = _make_placeholders()
tpu_computation = tpu.rewrite(model_fn, placeholders.keys())
sess.run(tpu.initialize_system())
sess.run(variables.global_variables_initializer())
result = sess.run(tpu_computation, placeholders)
sess.run(tpu.shutdown_system())
# TODO(b/36891278): supports non-flat returns lists in tpu.rewrite().
if len(result) == 1:
return result[0]
return result
elif device == "gpu":
with self.test_session(graph=ops.Graph(), use_gpu=True) as sess:
placeholders = _make_placeholders()
sess.run(variables.global_variables_initializer())
return sess.run(model_fn(placeholders.keys()), placeholders)
elif device == "cpu":
# TODO(power) -- will this interact poorly with cached GPU sessions?
with self.test_session(graph=ops.Graph(), use_gpu=False) as sess:
placeholders = _make_placeholders()
sess.run(variables.global_variables_initializer())
return sess.run(model_fn(placeholders.keys()), placeholders)
def _specialize_model(self, input_specs):
"""Specialize `self.model` (a Keras model) for the given input shapes."""
# Re-create our input and output layers inside our subgraph. They will be
# attached to the true computation when we clone our model in `tpu_fn`.
K.set_learning_phase(
self.execution_mode == model_fn_lib.ModeKeys.TRAIN)
# functools.partial and callable objects are not supported by tpu.rewrite
def _model_fn():
"""Compute fit/eval/predict for the TPU."""
is_training = self.execution_mode == model_fn_lib.ModeKeys.TRAIN
is_test = self.execution_mode == model_fn_lib.ModeKeys.EVAL
is_predict = self.execution_mode == model_fn_lib.ModeKeys.PREDICT
# During train/eval, we infeed our features as well as labels.
if is_training or is_test:
infeed_layers = self.model._input_layers + self.model._output_layers
else:
infeed_layers = self.model._input_layers
# Generate our infeed operation to read features & labels.
infeed_tensors = tpu_ops.infeed_dequeue_tuple(
dtypes=[spec.dtype for spec in input_specs],
shapes=[spec.shape for spec in input_specs],
name='infeed-%s' % self.execution_mode)
assert len(infeed_tensors) == len(infeed_layers), (
'Infeed inputs did not match model: %s vs %s',
(infeed_layers, infeed_tensors))
tpu_targets = []
tpu_inputs = []
# Sort infeed outputs into inputs and labels for calling our Keras model.
for tensor, layer in zip(infeed_tensors, infeed_layers):
if layer in self.model._input_layers:
tpu_inputs.append(
layers.Input(name=layer.name, tensor=tensor))
if layer in self.model._output_layers:
tpu_targets.append(tensor)
optimizer = self.model.optimizer
optimizer.iterations = training_util.get_or_create_global_step()
# Call our model with our infeed inputs (re-using the weights).
model_outputs = self.model(tpu_inputs)
child_model = models.Model(inputs=tpu_inputs,
outputs=model_outputs)
if is_training or is_test:
child_model.compile(
optimizer=self.model.optimizer,
loss=self.model.loss,
loss_weights=self.model.loss_weights,
metrics=self.model.metrics,
weighted_metrics=self.model.weighted_metrics,
target_tensors=tpu_targets,
)
# Compute our outfeed depending on the execution mode
if is_training:
child_model._make_train_function()
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype,
tensor.name)
for tensor in child_model.train_function.outputs
]
return [
child_model.train_function.updates_op,
tpu_ops.outfeed_enqueue_tuple(
child_model.train_function.outputs,
name='oufeed-enqueue-train')
]
elif is_test:
child_model._make_test_function()
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype,
tensor.name)
for tensor in child_model.test_function.outputs
]
return [
tpu_ops.outfeed_enqueue_tuple(
child_model.test_function.outputs,
name='outfeed-enqueue-test')
]
elif is_predict:
child_model._make_predict_function()
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype,
tensor.name)
for tensor in child_model.predict_function.outputs
]
return [
tpu_ops.outfeed_enqueue_tuple(
child_model.predict_function.outputs,
name='outfeed-enqueue-predict',
)
]
else:
assert False, 'Unexpected execution mode: %s' % self.execution_mode
# Capture outfeed metadata computed during the rewrite.
self._outfeed_spec = None
tpu_execute_op = tpu.rewrite(_model_fn)
K._initialize_variables(
K.get_session()) # pylint-disable: protected-access
# Generate CPU side operations to enqueue features/labels and dequeue
# outputs from the model call.
with ops.device('/device:TPU:0'):
infeed_tensors = []
for spec in input_specs:
infeed_tensors.append(
array_ops.placeholder(dtype=spec.dtype,
shape=spec.shape,
name='infeed-enqueue-%s' %
spec.name))
infeed_op = tpu_ops.infeed_enqueue_tuple(
infeed_tensors, [spec.shape for spec in input_specs],
name='infeed-enqueue-%s' % self.execution_mode)
outfeed_op = tpu_ops.outfeed_dequeue_tuple(
dtypes=[spec.dtype for spec in self._outfeed_spec],
shapes=[spec.shape for spec in self._outfeed_spec],
name='outfeed-dequeue-%s' % self.execution_mode)
return CompiledTPUOp(tpu_execute_op, infeed_tensors, infeed_op,
outfeed_op)
def _specialize_model(self, input_specs):
"""Specialize `self.model` (a Keras model) for the given input shapes."""
# Re-create our input and output layers inside our subgraph. They will be
# attached to the true computation when we clone our model in `tpu_fn`.
K.set_learning_phase(
self.execution_mode == model_fn_lib.ModeKeys.TRAIN
)
# functools.partial and callable objects are not supported by tpu.rewrite
def _model_fn():
"""Compute fit/eval/predict for the TPU."""
is_training = self.execution_mode == model_fn_lib.ModeKeys.TRAIN
is_test = self.execution_mode == model_fn_lib.ModeKeys.EVAL
is_predict = self.execution_mode == model_fn_lib.ModeKeys.PREDICT
# During train/eval, we infeed our features as well as labels.
if is_training or is_test:
infeed_layers = self.model._input_layers + self.model._output_layers
else:
infeed_layers = self.model._input_layers
# Generate our infeed operation to read features & labels.
infeed_tensors = tpu_ops.infeed_dequeue_tuple(
dtypes=[spec.dtype for spec in input_specs],
shapes=[spec.shape for spec in input_specs],
name='infeed-%s' % self.execution_mode)
assert len(infeed_tensors) == len(infeed_layers), (
'Infeed inputs did not match model: %s vs %s', (infeed_layers,
infeed_tensors))
tpu_targets = []
tpu_inputs = []
# Sort infeed outputs into inputs and labels for calling our Keras model.
for tensor, layer in zip(infeed_tensors, infeed_layers):
if layer in self.model._input_layers:
tpu_inputs.append(layers.Input(name=layer.name, tensor=tensor))
if layer in self.model._output_layers:
tpu_targets.append(tensor)
# Call our model with our infeed inputs (re-using the weights).
model_outputs = self.model(tpu_inputs)
child_model = models.Model(inputs=tpu_inputs, outputs=model_outputs)
if is_training or is_test:
child_model.compile(
optimizer=self.model.optimizer,
loss=self.model.loss,
loss_weights=self.model.loss_weights,
metrics=self.model.metrics,
weighted_metrics=self.model.weighted_metrics,
target_tensors=tpu_targets,
)
# Compute our outfeed depending on the execution mode
if is_training:
child_model._make_train_function()
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype, tensor.name)
for tensor in child_model.train_function.outputs
]
return [
child_model.train_function.updates_op,
tpu_ops.outfeed_enqueue_tuple(
child_model.train_function.outputs, name='oufeed-enqueue-train')
]
elif is_test:
child_model._make_test_function()
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype, tensor.name)
for tensor in child_model.test_function.outputs
]
return [
tpu_ops.outfeed_enqueue_tuple(
child_model.test_function.outputs, name='outfeed-enqueue-test')
]
elif is_predict:
child_model._make_predict_function()
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype, tensor.name)
for tensor in child_model.predict_function.outputs
]
return [
tpu_ops.outfeed_enqueue_tuple(
child_model.predict_function.outputs,
name='outfeed-enqueue-predict',
)
]
else:
assert False, 'Unexpected execution mode: %s' % self.execution_mode
# Capture outfeed metadata computed during the rewrite.
self._outfeed_spec = None
tpu_execute_op = tpu.rewrite(_model_fn)
# Generate CPU side operations to enqueue features/labels and dequeue
# outputs from the model call.
with ops.device('/device:TPU:0'):
infeed_tensors = []
for spec in input_specs:
infeed_tensors.append(
array_ops.placeholder(
dtype=spec.dtype,
shape=spec.shape,
name='infeed-enqueue-%s' % spec.name))
infeed_op = tpu_ops.infeed_enqueue_tuple(
infeed_tensors, [spec.shape for spec in input_specs],
name='infeed-enqueue-%s' % self.execution_mode)
outfeed_op = tpu_ops.outfeed_dequeue_tuple(
dtypes=[spec.dtype for spec in self._outfeed_spec],
shapes=[spec.shape for spec in self._outfeed_spec],
name='outfeed-dequeue-%s' % self.execution_mode)
return CompiledTPUOp(tpu_execute_op, infeed_tensors, infeed_op, outfeed_op)
def tpu_subgraph():
results = tpu.rewrite(functools.partial(model_fn, self.hparams),
args)
results = tf.reshape(results, [self.hparams.infer_batch_size, -1])
return self.vocab_table.lookup(tf.to_int64(results))
def tpu_subgraph():
return tpu.rewrite(functools.partial(model_fn, self.params),
inputs)