【hydra No.4】Adapt VIV to hydra

Author: GreatV

docs/zh/examples/viv.md

@@ -16,6 +16,13 @@
 
 ![VIV_1D_SpringDamper](https://paddle-org.bj.bcebos.com/paddlescience/docs/ViV/VIV_1D_SpringDamper.png)
 
+
+=== "模型训练命令"
+
+    ``` sh
+    python viv.py
+    ```
+
 ## 2. 问题定义
 
 本问题涉及的控制方程涉及三个物理量：$λ_1$、$λ_2$ 和 $ρ$，分别表示自然阻尼、结构特性刚度和质量，控制方程定义如下所示：
@@ -41,9 +48,9 @@ $$
 
 上式中 $g$ 即为 MLP 模型本身，用 PaddleScience 代码表示如下
 
-``` py linenums="28"
+``` py linenums="31"
 --8<--
-examples/fsi/viv.py:28:29
+examples/fsi/viv.py:31:32
 --8<--
 ```
 
@@ -56,9 +63,9 @@ examples/fsi/viv.py:28:29
 
 由于 VIV 使用的是 VIV 方程，因此可以直接使用 PaddleScience 内置的 `VIV`。
 
-``` py linenums="31"
+``` py linenums="34"
 --8<--
-examples/fsi/viv.py:31:32
+examples/fsi/viv.py:34:35
 --8<--
 ```
 
@@ -76,19 +83,19 @@ examples/fsi/viv.py:31:32
 
 在定义约束之前，需要给监督约束指定文件路径等数据读取配置。
 
-``` py linenums="34"
+``` py linenums="37"
 --8<--
-examples/fsi/viv.py:34:50
+examples/fsi/viv.py:37:52
 --8<--
 ```
 
 #### 3.4.1 监督约束
 
 由于我们以监督学习方式进行训练，此处采用监督约束 `SupervisedConstraint`：
 
-``` py linenums="51"
+``` py linenums="54"
 --8<--
-examples/fsi/viv.py:51:57
+examples/fsi/viv.py:54:60
 --8<--
 ```
 
@@ -102,29 +109,29 @@ examples/fsi/viv.py:51:57
 
 在监督约束构建完毕之后，以我们刚才的命名为关键字，封装到一个字典中，方便后续访问。
 
-``` py linenums="58"
+``` py linenums="61"
 --8<--
-examples/fsi/viv.py:58:61
+examples/fsi/viv.py:61:64
 --8<--
 ```
 
 ### 3.5 超参数设定
 
 接下来我们需要指定训练轮数和学习率，此处我们按实验经验，使用十万轮训练轮数，并每隔1000个epochs评估一次模型精度。
 
-``` py linenums="63"
+``` yaml linenums="39"
 --8<--
-examples/fsi/viv.py:63:65
+examples/fsi/conf/viv.yaml:39:40
 --8<--
 ```
 
 ### 3.6 优化器构建
 
 训练过程会调用优化器来更新模型参数，此处选择较为常用的 `Adam` 优化器和 `Step` 间隔衰减学习率。
 
-``` py linenums="67"
+``` py linenums="66"
 --8<--
-examples/fsi/viv.py:67:71
+examples/fsi/viv.py:66:68
 --8<--
 ```
 
@@ -136,9 +143,9 @@ examples/fsi/viv.py:67:71
 
 在训练过程中通常会按一定轮数间隔，用验证集（测试集）评估当前模型的训练情况，因此使用 `ppsci.validate.SupervisedValidator` 构建评估器。
 
-``` py linenums="73"
+``` py linenums="70"
 --8<--
-examples/fsi/viv.py:73:95
+examples/fsi/viv.py:70:92
 --8<--
 ```
 
@@ -152,19 +159,19 @@ examples/fsi/viv.py:73:95
 
 本文需要可视化的数据是 $t-\eta$ 和 $t-f$ 两组关系图，假设每个时刻 $t$ 的坐标是 $t_i$，则对应网络输出为 $\eta_i$，升力为 $f_i$，因此我们只需要将评估过程中产生的所有 $(t_i, \eta_i, f_i)$ 保存成图片即可。代码如下：
 
-``` py linenums="97"
+``` py linenums="94"
 --8<--
-examples/fsi/viv.py:97:116
+examples/fsi/viv.py:94:113
 --8<--
 ```
 
 ### 3.9 模型训练、评估与可视化
 
 完成上述设置之后，只需要将上述实例化的对象按顺序传递给 `ppsci.solver.Solver`，然后启动训练、评估、可视化。
 
-``` py linenums="118"
+``` py linenums="115"
 --8<--
-examples/fsi/viv.py:118:
+examples/fsi/viv.py:115:136
 --8<--
 ```
 

examples/fsi/conf/viv.yaml

@@ -0,0 +1,55 @@
+hydra:
+  run:
+    # dynamic output directory according to running time and override name
+    dir: outputs_VIV/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname}
+  job:
+    name: ${mode} # name of logfile
+    chdir: false # keep current working direcotry unchaned
+    config:
+      override_dirname:
+        exclude_keys:
+          - TRAIN.checkpoint_path
+          - TRAIN.pretrained_model_path
+          - EVAL.pretrained_model_path
+          - mode
+          - output_dir
+          - log_freq
+  sweep:
+    # output directory for multirun
+    dir: ${hydra.run.dir}
+    subdir: ./
+
+# general settings
+mode: train # running mode: train/eval
+seed: 42
+output_dir: ${hydra:run.dir}
+log_freq: 20
+
+
+# model settings
+MODEL:
+  input_keys: ["t_f"]
+  output_keys: ["eta"]
+  num_layers: 5
+  hidden_size: 50
+  activation: "tanh"
+
+# training settings
+TRAIN:
+  epochs: 100000
+  eval_freq: 1000
+  iters_per_epoch: 1
+  save_freq: 1
+  eval_during_train: true
+  optimizer:
+    epochs: 100000
+    iters_per_epoch: 1
+    learning_rate: 0.001
+    step_size: 20000
+    gamma: 0.9
+  pretrained_model_path: null
+  checkpoint_path: null
+
+# evaluation settings
+EVAL:
+  pretrained_model_path: null

examples/fsi/viv.py

@@ -12,27 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
+import hydra
+from omegaconf import DictConfig
+
 import ppsci
-from ppsci.utils import config
 from ppsci.utils import logger
 
-if __name__ == "__main__":
-    args = config.parse_args()
+
+def train(cfg: DictConfig):
     # set random seed for reproducibility
-    ppsci.utils.misc.set_random_seed(42)
+    ppsci.utils.misc.set_random_seed(cfg.seed)
+
     # set output directory
-    OUTPUT_DIR = "./output_viv" if args.output_dir is None else args.output_dir
-    # initialize logger
-    logger.init_logger("ppsci", f"{OUTPUT_DIR}/train.log", "info")
+    logger.init_logger("ppsci", os.path.join(cfg.output_dir, "train.log"), "info")
 
     # set model
-    model = ppsci.arch.MLP(("t_f",), ("eta",), 5, 50, "tanh")
+    model = ppsci.arch.MLP(**cfg.MODEL)
 
     # set equation
     equation = {"VIV": ppsci.equation.Vibration(2, -4, 0)}
 
     # set dataloader config
-    ITERS_PER_EPOCH = 1
     train_dataloader_cfg = {
         "dataset": {
             "name": "MatDataset",
@@ -48,6 +50,7 @@
             "shuffle": True,
         },
     }
+
     # set constraint
     sup_constraint = ppsci.constraint.SupervisedConstraint(
         train_dataloader_cfg,
@@ -60,14 +63,8 @@
         sup_constraint.name: sup_constraint,
     }
 
-    # set training hyper-parameters
-    EPOCHS = 100000 if args.epochs is None else args.epochs
-    EVAL_FREQ = 1000
-
     # set optimizer
-    lr_scheduler = ppsci.optimizer.lr_scheduler.Step(
-        EPOCHS, ITERS_PER_EPOCH, 0.001, step_size=20000, gamma=0.9
-    )()
+    lr_scheduler = ppsci.optimizer.lr_scheduler.Step(**cfg.TRAIN.optimizer)()
     optimizer = ppsci.optimizer.Adam(lr_scheduler)((model,) + tuple(equation.values()))
 
     # set validator
@@ -119,35 +116,110 @@
     solver = ppsci.solver.Solver(
         model,
         constraint,
-        OUTPUT_DIR,
+        cfg.output_dir,
         optimizer,
         lr_scheduler,
-        EPOCHS,
-        ITERS_PER_EPOCH,
+        cfg.TRAIN.epochs,
+        cfg.TRAIN.iters_per_epoch,
         eval_during_train=True,
-        eval_freq=EVAL_FREQ,
+        eval_freq=cfg.TRAIN.eval_freq,
         equation=equation,
         validator=validator,
         visualizer=visualizer,
     )
+
     # train model
     solver.train()
     # evaluate after finished training
     solver.eval()
     # visualize prediction after finished training
     solver.visualize()
 
-    # directly evaluate model from pretrained_model_path(optional)
-    logger.init_logger("ppsci", f"{OUTPUT_DIR}/eval.log", "info")
+
+def evaluate(cfg: DictConfig):
+    # set random seed for reproducibility
+    ppsci.utils.misc.set_random_seed(cfg.seed)
+
+    # set output directory
+    logger.init_logger("ppsci", os.path.join(cfg.output_dir, "eval.log"), "info")
+
+    # set model
+    model = ppsci.arch.MLP(**cfg.MODEL)
+
+    # set equation
+    equation = {"VIV": ppsci.equation.Vibration(2, -4, 0)}
+
+    # set validator
+    valid_dataloader_cfg = {
+        "dataset": {
+            "name": "MatDataset",
+            "file_path": "./VIV_Training_Neta100.mat",
+            "input_keys": ("t_f",),
+            "label_keys": ("eta", "f"),
+        },
+        "batch_size": 32,
+        "sampler": {
+            "name": "BatchSampler",
+            "drop_last": False,
+            "shuffle": False,
+        },
+    }
+    eta_mse_validator = ppsci.validate.SupervisedValidator(
+        valid_dataloader_cfg,
+        ppsci.loss.MSELoss("mean"),
+        {"eta": lambda out: out["eta"], **equation["VIV"].equations},
+        metric={"MSE": ppsci.metric.MSE()},
+        name="eta_mse",
+    )
+    validator = {eta_mse_validator.name: eta_mse_validator}
+
+    # set visualizer(optional)
+    visu_mat = ppsci.utils.reader.load_mat_file(
+        "./VIV_Training_Neta100.mat",
+        ("t_f", "eta_gt", "f_gt"),
+        alias_dict={"eta_gt": "eta", "f_gt": "f"},
+    )
+
+    visualizer = {
+        "visualize_u": ppsci.visualize.VisualizerScatter1D(
+            visu_mat,
+            ("t_f",),
+            {
+                r"$\eta$": lambda d: d["eta"],  # plot with latex title
+                r"$\eta_{gt}$": lambda d: d["eta_gt"],  # plot with latex title
+                r"$f$": equation["VIV"].equations["f"],  # plot with latex title
+                r"$f_{gt}$": lambda d: d["f_gt"],  # plot with latex title
+            },
+            num_timestamps=1,
+            prefix="viv_pred",
+        )
+    }
+
+    # initialize solver
     solver = ppsci.solver.Solver(
         model,
-        constraint,
-        OUTPUT_DIR,
+        output_dir=cfg.output_dir,
         equation=equation,
         validator=validator,
         visualizer=visualizer,
-        pretrained_model_path=f"{OUTPUT_DIR}/checkpoints/latest",
+        pretrained_model_path=cfg.EVAL.pretrained_model_path,
     )
+
+    # evaluate after finished training
     solver.eval()
-    # visualize prediction from pretrained_model_path(optional)
+    # visualize prediction after finished training
     solver.visualize()
+
+
+@hydra.main(version_base=None, config_path="./conf", config_name="viv.yaml")
+def main(cfg: DictConfig):
+    if cfg.mode == "train":
+        train(cfg)
+    elif cfg.mode == "eval":
+        evaluate(cfg)
+    else:
+        raise ValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'")
+
+
+if __name__ == "__main__":
+    main()