def cpu_model(self): cpu_model = models.Model( inputs=self.inputs, outputs=self.outputs, name=self.name, ) if self.optimizer: cpu_model.compile( optimizer=self.optimizer, loss=self.loss, metrics=self.metrics, loss_weights=self.loss_weights, ) return cpu_model
def cpu_model(self): return models.Model( inputs=self.inputs, outputs=self.outputs, name=self.name, )
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