【Hackathon 8th No.12】在 PaddleScience 中实现 SOAP 优化器 (#1102)

* Add SOAP optimizer

* update comments

* Add copyright

* Fix docs

* fix

* resolve reviewer issues

* add soap config

* fix

* fix

* resolve reviewer issues

Author: BeingGod

docs/zh/api/optimizer.md

@@ -11,5 +11,6 @@
         - OptimizerList
         - RMSProp
         - SGD
+        - SOAP
       show_root_heading: true
       heading_level: 3

docs/zh/examples/allen_cahn.md

@@ -9,7 +9,10 @@
     wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/AllenCahn/allen_cahn.mat -P ./dataset/
     # windows
     # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/AllenCahn/allen_cahn.mat --create-dirs -o ./dataset/allen_cahn.mat
+    # Using Adam optimizer
     python allen_cahn_piratenet.py
+    # Using SOAP optimizer
+    python allen_cahn_piratenet.py TRAIN.optim=soap TRAIN.lr_scheduler.warmup_epoch=5
     ```
 
 === "模型评估命令"
@@ -19,7 +22,10 @@
     wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/AllenCahn/allen_cahn.mat -P ./dataset/
     # windows
     # curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/AllenCahn/allen_cahn.mat --create-dirs -o ./dataset/allen_cahn.mat
+    # Using Adam pretrained model
     python allen_cahn_piratenet.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/AllenCahn/allen_cahn_piratenet_pretrained.pdparams
+    # Using SOAP pretrained model
+    python allen_cahn_piratenet.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/AllenCahn/allen_cahn_piratenet_soap_pretrained.pdparams
     ```
 
 === "模型导出命令"
@@ -41,6 +47,7 @@
 | 预训练模型  | 指标 |
 |:--| :--|
 | [allen_cahn_piratenet_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/AllenCahn/allen_cahn_piratenet_pretrained.pdparams) | L2Rel.u: 1.2e-05 |
+| [allen_cahn_piratenet_soap_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/AllenCahn/allen_cahn_piratenet_soap_pretrained.pdparams) | L2Rel.u: 6.8e-6 |
 
 ## 1. 背景简介
 

examples/allen_cahn/allen_cahn_piratenet.py

@@ -133,7 +133,15 @@ def gen_label_batch(input_batch):
     lr_scheduler = ppsci.optimizer.lr_scheduler.ExponentialDecay(
         **cfg.TRAIN.lr_scheduler
     )()
-    optimizer = ppsci.optimizer.Adam(lr_scheduler)(model)
+
+    if cfg.TRAIN.optim == "adam":
+        optimizer = ppsci.optimizer.Adam(lr_scheduler)(model)
+    elif cfg.TRAIN.optim == "soap":
+        optimizer = ppsci.optimizer.SOAP(lr_scheduler)(model)
+    else:
+        raise ValueError(
+            f"cfg.TRAIN.optim should be in ['adam','soap'], but got '{cfg.TRAIN.optim}'."
+        )
 
     # set validator
     tx_star = misc.cartesian_product(t_star, x_star).astype(dtype)

examples/allen_cahn/conf/allen_cahn_piratenet.yaml

@@ -54,13 +54,15 @@ TRAIN:
   save_freq: 10
   eval_during_train: true
   eval_freq: 10
+  optim: adam
   lr_scheduler:
     epochs: ${TRAIN.epochs}
     iters_per_epoch: ${TRAIN.iters_per_epoch}
     learning_rate: 1.0e-3
     gamma: 0.9
     decay_steps: 5000
     by_epoch: false
+    warmup_epoch: 0
   batch_size: 8192
   pretrained_model_path: null
   checkpoint_path: null

ppsci/optimizer/__init__.py

@@ -17,6 +17,7 @@
 from ppsci.optimizer import lr_scheduler
 from ppsci.optimizer.optimizer import LBFGS
 from ppsci.optimizer.optimizer import SGD
+from ppsci.optimizer.optimizer import SOAP
 from ppsci.optimizer.optimizer import Adam
 from ppsci.optimizer.optimizer import AdamW
 from ppsci.optimizer.optimizer import Momentum
@@ -32,6 +33,7 @@
     "RMSProp",
     "OptimizerList",
     "lr_scheduler",
+    "SOAP",
 ]
 
 

ppsci/optimizer/optimizer.py

@@ -27,6 +27,7 @@
 from paddle.incubate import optimizer as incubate_optim
 from typing_extensions import Literal
 
+from ppsci.optimizer.soap import SOAP as SOAP_impl
 from ppsci.utils import logger
 from ppsci.utils import misc
 
@@ -495,6 +496,104 @@ def _apply_decay_param_fun(self, name):
         return name not in self.no_weight_decay_param_name_list
 
 
+class SOAP:
+    """
+    Improving and Stabilizing Shampoo using Adam. Implements SOAP algorithm (https://arxiv.org/abs/2409.11321).
+
+    Args:
+        learning_rate (float, optional):
+            The learning rate to use. defaults to 0.003.
+        beta1 (float, optional):
+            Adam's betas parameters beta1. defaults to 0.95.
+        beta2 (float, optional):
+            Adam's betas parameters beta2. defaults to 0.95.
+        shampoo_beta (float, optional):
+            If >= 0, use this beta for the preconditioner (L and R in paper, state['GG'] below) moving average instead of betas[1].
+            defaults to -1.
+        epsilon (float, optional):
+            Adam's epsilon for numerical stability. defaults to 1e-08.
+        weight_decay (float, optional): weight decay coefficient. defaults to 0.01.
+        precondition_frequency (int, optional):
+            How often to update the preconditioner. defaults to 10.
+        max_precond_dim (int, optional):
+            Maximum dimension of the preconditioner.
+            Set to 10000, so that we exclude most common vocab sizes while including layers. defaults to 10000.
+        merge_dims (bool, optional):
+            Whether or not to merge dimensions of the preconditioner. defaults to `False`.
+        precondition_1d (bool, optional):
+            Whether or not to precondition 1D gradients. defaults to `False`.
+        normalize_grads (bool, optional):
+            Whether or not to normalize gradients per layer.
+            Helps at large precondition_frequency (~100 in our experiments),
+            but hurts performance at small precondition_frequency (~10 in our experiments). defaults to `False`.
+        data_format (str, optional):
+            Data format of the input for convolutional layers.
+            Should be "channels_last" for data_format of NHWC and "channels_first" for NCHW. defaults to `channels_first`.
+        correct_bias (bool, optional):
+            Whether or not to use bias correction in Adam. defaults to `True`.
+
+    Examples:
+        >>> import ppsci
+        >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20)
+        >>> opt = ppsci.optimizer.SOAP(1e-3)(model)
+    """
+
+    def __init__(
+        self,
+        learning_rate: float = 3e-3,
+        beta1: float = 0.95,
+        beta2: float = 0.95,
+        shampoo_beta: float = -1,
+        epsilon: float = 1e-8,
+        weight_decay: float = 0.01,
+        precondition_frequency: int = 10,
+        max_precond_dim: int = 10000,  #
+        merge_dims: bool = False,  # Merge dimensions till the product of the dimensions is less than or equal to max_precond_dim.
+        precondition_1d: bool = False,
+        normalize_grads: bool = False,
+        data_format: str = "channels_first",
+        correct_bias: bool = True,
+    ):
+        self.learning_rate = learning_rate
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.shampoo_beta = shampoo_beta
+        self.epsilon = epsilon
+        self.weight_decay = weight_decay
+        self.precondition_frequency = precondition_frequency
+        self.max_precond_dim = max_precond_dim
+        self.merge_dims = merge_dims
+        self.precondition_1d = precondition_1d
+        self.normalize_grads = normalize_grads
+        self.data_format = data_format
+        self.correct_bias = correct_bias
+
+    def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]):
+        # model_list is None in static graph
+        if not isinstance(model_list, (tuple, list)):
+            model_list = (model_list,)
+        parameters = (
+            sum([m.parameters() for m in model_list], []) if model_list else None
+        )
+        opt = SOAP_impl(
+            parameters=parameters,
+            learning_rate=self.learning_rate,
+            beta1=self.beta1,
+            beta2=self.beta2,
+            shampoo_beta=self.shampoo_beta,
+            epsilon=self.epsilon,
+            weight_decay=self.weight_decay,
+            precondition_frequency=self.precondition_frequency,
+            max_precond_dim=self.max_precond_dim,
+            merge_dims=self.merge_dims,
+            precondition_1d=self.precondition_1d,
+            normalize_grads=self.normalize_grads,
+            data_format=self.data_format,
+            correct_bias=self.correct_bias,
+        )
+        return opt
+
+
 class OptimizerList:
     """OptimizerList which wrap more than one optimizer.
     NOTE: LBFGS is not supported yet.