Module: tfm.optimization

View source on GitHub

Optimization package definition.

Modules

adafactor_optimizer module: Adafactor optimizer.

base_config module: Base configurations to standardize experiments.

configs module

ema_optimizer module: Exponential moving average optimizer.

lamb module: Layer-wise Adaptive Moments (LAMB) optimizer.

lars module: Layer-wise adaptive rate scaling optimizer.

legacy_adamw module: Adam optimizer with weight decay that exactly matches the original BERT.

lr_cfg module: Dataclasses for learning rate schedule config.

lr_schedule module: Learning rate schedule classes.

math module: This module provides access to the mathematical functions defined by the C standard.

oneof module: Config class that supports oneof functionality.

opt_cfg module: Dataclasses for optimizer configs.

optimizer_factory module: Optimizer factory class.

slide_optimizer module: SLIDE optimizer.

Classes

class AdafactorConfig: Configuration for Adafactor optimizer.

class AdagradConfig: Configuration for Adagrad optimizer.

class AdamConfig: Configuration for Adam optimizer.

class AdamExperimentalConfig: Configuration for experimental Adam optimizer.

class AdamWeightDecayConfig: Configuration for Adam optimizer with weight decay.

class AdamWeightDecayExperimentalConfig: Configuration for Adam optimizer with weight decay.

class BaseOptimizerConfig: Base optimizer config.

class ConstantLrConfig: Configuration for constant learning rate.

class CosineDecayWithOffset: A LearningRateSchedule that uses a cosine decay with optional warmup.

class CosineLrConfig: Configuration for Cosine learning rate decay.

class DirectPowerDecay: Learning rate schedule follows lr * (step)^power.

class DirectPowerLrConfig: Configuration for DirectPower learning rate decay.

class EMAConfig: Exponential moving average optimizer config.

class ExponentialDecayWithOffset: A LearningRateSchedule that uses an exponential decay schedule.

class ExponentialLrConfig: Configuration for exponential learning rate decay.

class ExponentialMovingAverage: Optimizer that computes an exponential moving average of the variables.

class LAMBConfig: Configuration for LAMB optimizer.

class LARSConfig: Layer-wise adaptive rate scaling config.

class LinearWarmup: Linear warmup schedule.

class LinearWarmupConfig: Configuration for linear warmup schedule config.

class LrConfig: Configuration for lr schedule.

class OptimizationConfig: Configuration for optimizer and learning rate schedule.

class OptimizerConfig: Configuration for optimizer.

class OptimizerFactory: Optimizer factory class.

class PiecewiseConstantDecayWithOffset: A LearningRateSchedule that uses a piecewise constant decay schedule.

class PolynomialDecayWithOffset: A LearningRateSchedule that uses a polynomial decay schedule.

class PolynomialLrConfig: Configuration for polynomial learning rate decay.

class PolynomialWarmUp: Applies polynomial warmup schedule on a given learning rate decay schedule.

class PolynomialWarmupConfig: Configuration for linear warmup schedule config.

class PowerAndLinearDecay: Learning rate schedule with multiplied by linear decay at the end.

class PowerAndLinearDecayLrConfig: Configuration for DirectPower learning rate decay.

class PowerDecayWithOffset: Power learning rate decay with offset.

class PowerDecayWithOffsetLrConfig: Configuration for power learning rate decay with step offset.

class RMSPropConfig: Configuration for RMSProp optimizer.

class SGDConfig: Configuration for SGD optimizer.

class SGDExperimentalConfig: Configuration for SGD optimizer.

class SLIDEConfig: Configuration for SLIDE optimizer.

class StepCosineDecayWithOffset: Stepwise cosine learning rate decay with offset.

class StepCosineLrConfig: Configuration for stepwise learning rate decay.

class StepwiseLrConfig: Configuration for stepwise learning rate decay.

class WarmupConfig: Configuration for lr schedule.

Functions

register_optimizer_cls(...): Register customize optimizer cls.