tf.tpu.experimental.embedding.Adam

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

tf.tpu.experimental.embedding.Adam(
    learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-07, lazy_adam=True,
    sum_inside_sqrt=True, use_gradient_accumulation=True, clip_weight_min=None,
    clip_weight_max=None, weight_decay_factor=None,
    multiply_weight_decay_factor_by_learning_rate=None,
    slot_variable_creation_fn=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» для полного описания этих параметров и их влияния на алгоритм оптимизации.