tf.tpu.experimental.embedding.Adam

Параметры оптимизации для Adam с TPU встраиванием.

tf.tpu.experimental.embedding.Adam(
    learning_rate: Union[float, Callable[[], float]] = 0.001,
    beta_1: float = 0.9,
    beta_2: float = 0.999,
    epsilon: float = 1e-07,
    lazy_adam: bool = True,
    sum_inside_sqrt: bool = True,
    use_gradient_accumulation: bool = True,
    clip_weight_min: Optional[float] = None,
    clip_weight_max: Optional[float] = None,
    weight_decay_factor: Optional[float] = None,
    multiply_weight_decay_factor_by_learning_rate: bool = None,
    slot_variable_creation_fn: Optional[SlotVarCreationFnType] = None,
    clipvalue: Optional[ClipValueType] = None
)

Передайте это в tf.tpu.experimental.embedding.TPUEmbedding через аргумент optimizer, чтобы установить глобальный оптимизатор и его параметры:

Примечание: По умолчанию этот оптимизатор ленивый, т. е. он не будет применять обновление градиента нулем к строкам, которые не были просмотрены. Вы можете изменить это поведение, задав lazy_adam значение False.

embedding = tf.tpu.experimental.embedding.TPUEmbedding(
    ...
    optimizer=tf.tpu.experimental.embedding.Adam(0.1))

Его также можно использовать в tf.tpu.experimental.embedding.TableConfig в качестве параметра оптимизатора, чтобы установить оптимизатор, специфичный для таблицы. Это переопределит оптимизатор и параметры глобального оптимизатора встраивания, определенного выше:

table_one = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...,
    optimizer=tf.tpu.experimental.embedding.Adam(0.2))
table_two = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...)

feature_config = (
    tf.tpu.experimental.embedding.FeatureConfig(
        table=table_one),
    tf.tpu.experimental.embedding.FeatureConfig(
        table=table_two))

embedding = tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config=feature_config,
    batch_size=...
    optimizer=tf.tpu.experimental.embedding.Adam(0.1))

В приведенном выше примере первое признак будет находиться в таблице, имеющей скорость обучения 0,2, а второе — в таблице, имеющей скорость обучения 0,1.

Полное описание этих параметров и их влияние на алгоритм оптимизатора см. в файле 'tensorflow/core/protobuf/tpu/optimization_parameters.proto'.