TPU compatible Keras model wrapper.
| 1376 | |
| 1377 | |
| 1378 | class KerasTPUModel(models.Model): |
| 1379 | """TPU compatible Keras model wrapper.""" |
| 1380 | |
| 1381 | def __init__(self, cpu_model, strategy): |
| 1382 | super(models.Model, self).__init__( # pylint: disable=bad-super-call |
| 1383 | inputs=cpu_model.inputs, |
| 1384 | outputs=cpu_model.outputs, |
| 1385 | name=cpu_model.name, |
| 1386 | ) |
| 1387 | if tf2.enabled(): |
| 1388 | raise RuntimeError( |
| 1389 | 'Keras support is now deprecated in support of TPU Strategy. ' |
| 1390 | 'Please follow the distribution strategy guide on tensorflow.org ' |
| 1391 | 'to migrate to the 2.0 supported version.') |
| 1392 | else: |
| 1393 | logging.warning( |
| 1394 | 'Keras support is now deprecated in support of TPU Strategy. ' |
| 1395 | 'Please follow the distribution strategy guide on tensorflow.org ' |
| 1396 | 'to migrate to the 2.0 supported version.') |
| 1397 | # Create a mapping from numpy arrays to infeed managers. |
| 1398 | # Note: uses a list of tuples instead of a map because numpy arrays are |
| 1399 | # not hashable. |
| 1400 | self._numpy_to_infeed_manager_list = [] |
| 1401 | |
| 1402 | # Add distribution specific arguments since we don't call the Model init. |
| 1403 | self._distribution_strategy = None |
| 1404 | self._compile_distribution = None |
| 1405 | |
| 1406 | self.predict_function = None |
| 1407 | self.test_function = None |
| 1408 | self.train_function = None |
| 1409 | self._stateful_metric_functions = [] |
| 1410 | |
| 1411 | cluster_resolver = strategy._tpu_cluster_resolver |
| 1412 | self._tpu_name_or_address = cluster_resolver.get_master() |
| 1413 | self._cpu_model = cpu_model |
| 1414 | self._tpu_assignment = strategy._make_assignment_for_model(cpu_model) |
| 1415 | self._tpu_model = None |
| 1416 | self._tpu_weights_initialized = False |
| 1417 | |
| 1418 | # If the input CPU model has already been compiled, compile our TPU model |
| 1419 | # immediately. |
| 1420 | if self._cpu_model.optimizer: |
| 1421 | self.compile( |
| 1422 | self._cpu_model.optimizer, |
| 1423 | self._cpu_model.loss, |
| 1424 | self._cpu_model._compile_metrics, |
| 1425 | self._cpu_model.loss_weights, |
| 1426 | self._cpu_model.sample_weight_mode, |
| 1427 | self._cpu_model._compile_weighted_metrics, |
| 1428 | self._cpu_model.target_tensors, |
| 1429 | ) |
| 1430 | |
| 1431 | # This flag must be disabled upon model mutation, such as changing the model |
| 1432 | # layers or recompiling the model to use a different optimizer. New function |
| 1433 | # definitions are generated whenever this flag is disabled, ensuring that |
| 1434 | # internal graph functions are always using the current model structure. |
| 1435 | # |