API for using TPU for embedding lookups. Args: table_to_config_dict: A dictionary mapping from string of table name to `TableConfig`. Table refers to an embedding table, e.g. `params` argument to `tf.nn.embedding_lookup_sparse()`. feature_to_config_dict: A dictionary
(self,
table_to_config_dict,
feature_to_config_dict,
batch_size,
mode,
master=None,
optimization_parameters=None,
cluster_def=None,
pipeline_execution_with_tensor_core=False,
partition_strategy='div',
device_config=None)
| 506 | # we can add `optimization_parameters` to `TableConfig` to override this |
| 507 | # global setting. |
| 508 | def __init__(self, |
| 509 | table_to_config_dict, |
| 510 | feature_to_config_dict, |
| 511 | batch_size, |
| 512 | mode, |
| 513 | master=None, |
| 514 | optimization_parameters=None, |
| 515 | cluster_def=None, |
| 516 | pipeline_execution_with_tensor_core=False, |
| 517 | partition_strategy='div', |
| 518 | device_config=None): |
| 519 | """API for using TPU for embedding lookups. |
| 520 | |
| 521 | Args: |
| 522 | table_to_config_dict: A dictionary mapping from string of table name to |
| 523 | `TableConfig`. Table refers to an embedding table, e.g. `params` |
| 524 | argument to `tf.nn.embedding_lookup_sparse()`. |
| 525 | feature_to_config_dict: A dictionary mapping from string of feature name |
| 526 | to `FeatureConfig`. Feature refers to ids to lookup in embedding table, |
| 527 | e.g. `sp_ids` argument to `tf.nn.embedding_lookup_sparse()`. |
| 528 | batch_size: An `int` representing the global batch size. |
| 529 | mode: `TRAINING` or `INFERENCE`. |
| 530 | master: A `string` representing the TensorFlow master to use. |
| 531 | optimization_parameters: `AdagradParameters`, `AdamParameters`, |
| 532 | `Stochasticgradientdescentparameters`. Must be set in training unless |
| 533 | all tables specify their own optimizers. And it must be `None` in |
| 534 | inference. |
| 535 | cluster_def: A ClusterDef object describing the TPU cluster. |
| 536 | pipeline_execution_with_tensor_core: setting this to `True` makes training |
| 537 | faster, but trained model will be different if step N and step N+1 |
| 538 | involve the same set of embedding IDs. Please see |
| 539 | `tpu_embedding_configuration.proto` for details. |
| 540 | partition_strategy: A string, either 'mod' or 'div', specifying how to map |
| 541 | the lookup id to the embedding tensor. For more information see |
| 542 | `tf.nn.embedding_lookup_sparse`. |
| 543 | device_config: A DeviceConfig instance, used when `master` and |
| 544 | `cluster_def` are both `None`. |
| 545 | |
| 546 | Raises: |
| 547 | ValueError: if any input is invalid. |
| 548 | """ |
| 549 | if partition_strategy not in ('div', 'mod'): |
| 550 | raise ValueError( |
| 551 | 'Invalid partition_strategy {}'.format(partition_strategy)) |
| 552 | self._partition_strategy = partition_strategy |
| 553 | |
| 554 | _validate_table_to_config_dict(table_to_config_dict) |
| 555 | # Avoid nondeterminism from `Dict` iteration order by using `OrderedDict`. |
| 556 | self._table_to_config_dict = _create_ordered_dict(table_to_config_dict) |
| 557 | |
| 558 | _validate_feature_to_config_dict(table_to_config_dict, |
| 559 | feature_to_config_dict) |
| 560 | self._feature_to_config_dict = _create_ordered_dict(feature_to_config_dict) |
| 561 | self._table_to_features_dict, self._table_to_num_features_dict = ( |
| 562 | _create_table_to_features_and_num_features_dicts( |
| 563 | self._feature_to_config_dict)) |
| 564 | self._combiners = _create_combiners(self._table_to_config_dict, |
| 565 | self._table_to_features_dict) |
no test coverage detected