|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """Task for the Ranking model."""
|
|
|
| import math
|
| from typing import Dict, List, Optional, Union
|
|
|
| import tensorflow as tf, tf_keras
|
| import tensorflow_recommenders as tfrs
|
|
|
| from official.core import base_task
|
| from official.core import config_definitions
|
| from official.recommendation.ranking import common
|
| from official.recommendation.ranking.configs import config
|
| from official.recommendation.ranking.data import data_pipeline
|
|
|
| RuntimeConfig = config_definitions.RuntimeConfig
|
|
|
|
|
| def _get_tpu_embedding_feature_config(
|
| vocab_sizes: List[int],
|
| embedding_dim: Union[int, List[int]],
|
| table_name_prefix: str = 'embedding_table',
|
| batch_size: Optional[int] = None
|
| ) -> Dict[str, tf.tpu.experimental.embedding.FeatureConfig]:
|
| """Returns TPU embedding feature config.
|
|
|
| i'th table config will have vocab size of vocab_sizes[i] and embedding
|
| dimension of embedding_dim if embedding_dim is an int or embedding_dim[i] if
|
| embedding_dim is a list).
|
| Args:
|
| vocab_sizes: List of sizes of categories/id's in the table.
|
| embedding_dim: An integer or a list of embedding table dimensions.
|
| table_name_prefix: a prefix for embedding tables.
|
| batch_size: Per-replica batch size.
|
| Returns:
|
| A dictionary of feature_name, FeatureConfig pairs.
|
| """
|
| if isinstance(embedding_dim, List):
|
| if len(vocab_sizes) != len(embedding_dim):
|
| raise ValueError(
|
| f'length of vocab_sizes: {len(vocab_sizes)} is not equal to the '
|
| f'length of embedding_dim: {len(embedding_dim)}')
|
| elif isinstance(embedding_dim, int):
|
| embedding_dim = [embedding_dim] * len(vocab_sizes)
|
| else:
|
| raise ValueError('embedding_dim is not either a list or an int, got '
|
| f'{type(embedding_dim)}')
|
|
|
| feature_config = {}
|
|
|
| for i, vocab_size in enumerate(vocab_sizes):
|
| table_config = tf.tpu.experimental.embedding.TableConfig(
|
| vocabulary_size=vocab_size,
|
| dim=embedding_dim[i],
|
| combiner='mean',
|
| initializer=tf.initializers.TruncatedNormal(
|
| mean=0.0, stddev=1 / math.sqrt(embedding_dim[i])),
|
| name=table_name_prefix + '_%02d' % i)
|
| feature_config[str(i)] = tf.tpu.experimental.embedding.FeatureConfig(
|
| name=str(i),
|
| table=table_config,
|
| output_shape=[batch_size] if batch_size else None,
|
| )
|
|
|
| return feature_config
|
|
|
|
|
| class RankingTask(base_task.Task):
|
| """A task for Ranking Model."""
|
|
|
| def __init__(self,
|
| params: config.Task,
|
| trainer_config: config.TrainerConfig,
|
| logging_dir: Optional[str] = None,
|
| steps_per_execution: int = 1,
|
| name: Optional[str] = None):
|
| """Task initialization.
|
|
|
| Args:
|
| params: the RankingModel task configuration instance.
|
| trainer_config: Trainer configuration instance.
|
| logging_dir: a string pointing to where the model, summaries etc. will be
|
| saved.
|
| steps_per_execution: Int. Defaults to 1. The number of batches to run
|
| during each `tf.function` call. It's used for compile/fit API.
|
| name: the task name.
|
| """
|
| super().__init__(params, logging_dir, name=name)
|
| self._trainer_config = trainer_config
|
| self._optimizer_config = trainer_config.optimizer_config
|
| self._steps_per_execution = steps_per_execution
|
|
|
| def build_inputs(self, params, input_context=None):
|
| """Builds classification input."""
|
|
|
| dataset = data_pipeline.CriteoTsvReader(
|
| file_pattern=params.input_path,
|
| params=params,
|
| vocab_sizes=self.task_config.model.vocab_sizes,
|
| num_dense_features=self.task_config.model.num_dense_features,
|
| use_synthetic_data=self.task_config.use_synthetic_data)
|
|
|
| return dataset(input_context)
|
|
|
| @classmethod
|
| def create_optimizer(cls, optimizer_config: config.OptimizationConfig,
|
| runtime_config: Optional[RuntimeConfig] = None) -> None:
|
| """See base class. Return None, optimizer is set in `build_model`."""
|
| return None
|
|
|
| def build_model(self) -> tf_keras.Model:
|
| """Creates Ranking model architecture and Optimizers.
|
|
|
| The RankingModel uses different optimizers/learning rates for embedding
|
| variables and dense variables.
|
|
|
| Returns:
|
| A Ranking model instance.
|
| """
|
| lr_config = self.optimizer_config.lr_config
|
| lr_callable = common.WarmUpAndPolyDecay(
|
| batch_size=self.task_config.train_data.global_batch_size,
|
| decay_exp=lr_config.decay_exp,
|
| learning_rate=lr_config.learning_rate,
|
| warmup_steps=lr_config.warmup_steps,
|
| decay_steps=lr_config.decay_steps,
|
| decay_start_steps=lr_config.decay_start_steps)
|
| embedding_optimizer = tf_keras.optimizers.get(
|
| self.optimizer_config.embedding_optimizer, use_legacy_optimizer=True)
|
| embedding_optimizer.learning_rate = lr_callable
|
|
|
| dense_optimizer = tf_keras.optimizers.get(
|
| self.optimizer_config.dense_optimizer, use_legacy_optimizer=True)
|
| if self.optimizer_config.dense_optimizer == 'SGD':
|
| dense_lr_config = self.optimizer_config.dense_sgd_config
|
| dense_lr_callable = common.WarmUpAndPolyDecay(
|
| batch_size=self.task_config.train_data.global_batch_size,
|
| decay_exp=dense_lr_config.decay_exp,
|
| learning_rate=dense_lr_config.learning_rate,
|
| warmup_steps=dense_lr_config.warmup_steps,
|
| decay_steps=dense_lr_config.decay_steps,
|
| decay_start_steps=dense_lr_config.decay_start_steps)
|
| dense_optimizer.learning_rate = dense_lr_callable
|
|
|
| feature_config = _get_tpu_embedding_feature_config(
|
| embedding_dim=self.task_config.model.embedding_dim,
|
| vocab_sizes=self.task_config.model.vocab_sizes,
|
| batch_size=self.task_config.train_data.global_batch_size
|
| // tf.distribute.get_strategy().num_replicas_in_sync,
|
| )
|
|
|
| embedding_layer = tfrs.experimental.layers.embedding.PartialTPUEmbedding(
|
| feature_config=feature_config,
|
| optimizer=embedding_optimizer,
|
| pipeline_execution_with_tensor_core=self.trainer_config.pipeline_sparse_and_dense_execution,
|
| size_threshold=self.task_config.model.size_threshold,
|
| )
|
|
|
| if self.task_config.model.interaction == 'dot':
|
| feature_interaction = tfrs.layers.feature_interaction.DotInteraction(
|
| skip_gather=True)
|
| elif self.task_config.model.interaction == 'cross':
|
| feature_interaction = tf_keras.Sequential([
|
| tf_keras.layers.Concatenate(),
|
| tfrs.layers.feature_interaction.Cross()
|
| ])
|
| else:
|
| raise ValueError(
|
| f'params.task.model.interaction {self.task_config.model.interaction} '
|
| f'is not supported it must be either \'dot\' or \'cross\'.')
|
|
|
| model = tfrs.experimental.models.Ranking(
|
| embedding_layer=embedding_layer,
|
| bottom_stack=tfrs.layers.blocks.MLP(
|
| units=self.task_config.model.bottom_mlp, final_activation='relu'),
|
| feature_interaction=feature_interaction,
|
| top_stack=tfrs.layers.blocks.MLP(
|
| units=self.task_config.model.top_mlp, final_activation='sigmoid'),
|
| )
|
| optimizer = tfrs.experimental.optimizers.CompositeOptimizer([
|
| (embedding_optimizer, lambda: model.embedding_trainable_variables),
|
| (dense_optimizer, lambda: model.dense_trainable_variables),
|
| ])
|
|
|
| model.compile(optimizer, steps_per_execution=self._steps_per_execution)
|
| return model
|
|
|
| def train_step(
|
| self,
|
| inputs: Dict[str, tf.Tensor],
|
| model: tf_keras.Model,
|
| optimizer: tf_keras.optimizers.Optimizer,
|
| metrics: Optional[List[tf_keras.metrics.Metric]] = None) -> tf.Tensor:
|
| """See base class."""
|
|
|
| assert metrics == model.metrics
|
| return model.train_step(inputs)
|
|
|
| def validation_step(
|
| self,
|
| inputs: Dict[str, tf.Tensor],
|
| model: tf_keras.Model,
|
| metrics: Optional[List[tf_keras.metrics.Metric]] = None) -> tf.Tensor:
|
| """See base class."""
|
|
|
| assert metrics == model.metrics
|
| return model.test_step(inputs)
|
|
|
| @property
|
| def trainer_config(self) -> config.TrainerConfig:
|
| return self._trainer_config
|
|
|
| @property
|
| def optimizer_config(self) -> config.OptimizationConfig:
|
| return self._optimizer_config
|
|
|
