tfm.optimization.register_optimizer_cls

Register customize optimizer cls.

View aliases

Main aliases

tfm.optimization.optimizer_factory.register_optimizer_cls

tfm.optimization.register_optimizer_cls(
    key: str,
    optimizer_config_cls: Union[tf.keras.optimizers.Optimizer, tf.keras.optimizers.legacy.Optimizer,
        tf.keras.optimizers.experimental.Optimizer],
    use_legacy_optimizer: bool = True
)

The user will still need to subclass data classes in configs.optimization_config to be used with OptimizerFactory.

Args
key	A string to that the optimizer_config_cls is registered with.
optimizer_config_cls	A class which inherits tf.keras.optimizers.Optimizer.
use_legacy_optimizer	A boolean that indicates if using legacy optimizers.