Add weight decay regex to Adafactor and expose weight decay params to its config.

PiperOrigin-RevId: 448278992

Add weight decay regex to Adafactor and expose weight decay params to its config.
PiperOrigin-RevId: 448278992
fbc6a068 · A. Unique TensorFlower · 1f32d4e4 · fbc6a068
Commit fbc6a068 authored May 12, 2022 by A. Unique TensorFlower
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

official/modeling/optimization/configs/optimizer_config.py official/modeling/optimization/configs/optimizer_config.py +2 -0

No files found.
--- a/official/modeling/optimization/configs/optimizer_config.py
+++ b/official/modeling/optimization/configs/optimizer_config.py
@@ -311,3 +311,5 @@ class AdafactorConfig(BaseOptimizerConfig):
  min_dim_size_to_factor: int = 128
  epsilon1: float = 1e-30
  epsilon2: float = 1e-3
+  weight_decay: Optional[float] = None
+  include_in_weight_decay: Optional[str] = None