【hydra No.19】改造 phylstm

docs/zh/examples/phylstm.md

@@ -29,9 +29,9 @@ $$
 
 在 PhyLSTM 问题中，建立 LSTM 网络 Deep LSTM network，用 PaddleScience 代码表示如下
 
-``` py linenums="105"
+``` py linenums="102"
 --8<--
-examples/phylstm/phylstm2.py:105:106
+examples/phylstm/phylstm2.py:102:107
 --8<--
 ```
 
@@ -47,39 +47,39 @@ wget -P ./ https://paddle-org.bj.bcebos.com/paddlescience/datasets/PhyLSTM/data_
 
 本案例涉及读取数据构建，如下所示
 
-``` py linenums="38"
+``` py linenums="37"
 --8<--
-examples/phylstm/phylstm2.py:38:105
+examples/phylstm/phylstm2.py:37:100
 --8<--
 ```
 
 ### 3.3 约束构建
 
 设置训练数据集和损失计算函数，返回字段，代码如下所示：
 
-``` py linenums="118"
+``` py linenums="119"
 --8<--
-examples/phylstm/phylstm2.py:118:136
+examples/phylstm/phylstm2.py:119:137
 --8<--
 ```
 
 ### 3.4 评估器构建
 
 设置评估数据集和损失计算函数，返回字段，代码如下所示：
 
-``` py linenums="139"
+``` py linenums="140"
 --8<--
-examples/phylstm/phylstm2.py:139:158
+examples/phylstm/phylstm2.py:140:159
 --8<--
 ```
 
 ### 3.5 超参数设定
 
 接下来我们需要指定训练轮数，此处我们按实验经验，使用 100 轮训练轮数。
 
-``` py linenums="36"
+``` py linenums="170"
 --8<--
-examples/phylstm/phylstm2.py:36:36
+examples/phylstm/phylstm2.py:170:170
 --8<--
 ```
 
@@ -99,15 +99,15 @@ examples/phylstm/phylstm2.py:163:163
 
 ``` py linenums="164"
 --8<--
-examples/phylstm/phylstm2.py:164:175
+examples/phylstm/phylstm2.py:164:178
 --8<--
 ```
 
 最后启动训练、评估即可：
 
-``` py linenums="177"
+``` py linenums="180"
 --8<--
-examples/phylstm/phylstm2.py:177:180
+examples/phylstm/phylstm2.py:180:183
 --8<--
 ```
 

examples/phylstm/conf/phylstm2.yaml

@@ -0,0 +1,39 @@
+hydra:
+  run:
+    # dynamic output directory according to running time and override name
+    dir: outputs_PhyLSTM2/${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
+  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
+data_file: data_boucwen.mat
+
+# model settings
+MODEL:
+  phylstm2_net:
+    input_size: 1
+    hidden_size: 100
+    model_type: 2
+
+# training settings
+TRAIN:
+  epochs: 100
+  iters_per_epoch: 1
+  save_freq: 50
+  pretrained_model_path: null
+  checkpoint_path: null
+  learning_rate: 0.001
+
+# evaluation settings
+EVAL:
+  pretrained_model_path: null
+  eval_with_no_grad: true

examples/phylstm/conf/phylstm3.yaml

@@ -0,0 +1,39 @@
+hydra:
+  run:
+    # dynamic output directory according to running time and override name
+    dir: outputs_PhyLSTM3/${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
+  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
+data_file: data_boucwen.mat
+
+# model settings
+MODEL:
+  phylstm3_net:
+    input_size: 1
+    hidden_size: 100
+    model_type: 3
+
+# training settings
+TRAIN:
+  epochs: 200
+  iters_per_epoch: 1
+  save_freq: 50
+  pretrained_model_path: null
+  checkpoint_path: null
+  learning_rate: 0.001
+
+# evaluation settings
+EVAL:
+  pretrained_model_path: null
+  eval_with_no_grad: true

examples/phylstm/phylstm2.py

@@ -16,26 +16,25 @@
 Reference: https://github.com/zhry10/PhyLSTM.git
 """
 
+from os import path as osp
+
 import functions
+import hydra
 import numpy as np
 import scipy.io
+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)
-    # set output directory
-    OUTPUT_DIR = "./output_PhyLSTM2" if not args.output_dir else args.output_dir
+    ppsci.utils.misc.set_random_seed(cfg.seed)
     # initialize logger
-    logger.init_logger("ppsci", f"{OUTPUT_DIR}/train.log", "info")
-    # set training hyper-parameters
-    EPOCHS = 100 if not args.epochs else args.epochs
+    logger.init_logger("ppsci", osp.join(cfg.output_dir, "train.log"), "info")
 
-    mat = scipy.io.loadmat("data_boucwen.mat")
+    mat = scipy.io.loadmat(cfg.data_file)
     ag_data = mat["input_tf"]  # ag, ad, av
     u_data = mat["target_X_tf"]
     ut_data = mat["target_Xd_tf"]
@@ -90,9 +89,6 @@
     eta_tt_star = u_tt_all[0:10]
     ag_c_star = ag_all[0:50]
     lift_star = -ag_c_star
-    eta_c_star = u_all[0:50]
-    eta_t_c_star = u_t_all[0:50]
-    eta_tt_c_star = u_tt_all[0:50]
 
     eta = eta_star
     ag = ag_star
@@ -103,7 +99,12 @@
     g = -eta_tt - ag
     phi_t = np.repeat(phi_t0, ag_c_star.shape[0], axis=0)
 
-    model = ppsci.arch.DeepPhyLSTM(1, eta.shape[2], 100, 2)
+    model = ppsci.arch.DeepPhyLSTM(
+        cfg.MODEL.phylstm2_net.input_size,
+        eta.shape[2],
+        cfg.MODEL.phylstm2_net.hidden_size,
+        cfg.MODEL.phylstm2_net.model_type,
+    )
     model.register_input_transform(functions.transform_in)
     model.register_output_transform(functions.transform_out)
 
@@ -113,7 +114,7 @@
         label_dict_train,
         input_dict_val,
         label_dict_val,
-    ) = dataset_obj.get(EPOCHS)
+    ) = dataset_obj.get(cfg.TRAIN.epochs)
 
     sup_constraint_pde = ppsci.constraint.SupervisedConstraint(
         {
@@ -159,22 +160,161 @@
     validator_pde = {sup_validator_pde.name: sup_validator_pde}
 
     # initialize solver
-    ITERS_PER_EPOCH = 1
-    optimizer = ppsci.optimizer.Adam(1e-3)(model)
+    optimizer = ppsci.optimizer.Adam(cfg.TRAIN.learning_rate)(model)
     solver = ppsci.solver.Solver(
         model,
         constraint_pde,
-        OUTPUT_DIR,
+        cfg.output_dir,
         optimizer,
         None,
-        EPOCHS,
-        ITERS_PER_EPOCH,
-        save_freq=50,
+        cfg.TRAIN.epochs,
+        cfg.TRAIN.iters_per_epoch,
+        save_freq=cfg.TRAIN.save_freq,
+        log_freq=cfg.log_freq,
+        seed=cfg.seed,
         validator=validator_pde,
-        eval_with_no_grad=True,
+        checkpoint_path=cfg.TRAIN.checkpoint_path,
+        eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
     )
 
     # train model
     solver.train()
     # evaluate after finished training
     solver.eval()
+
+
+def evaluate(cfg: DictConfig):
+    # set random seed for reproducibility
+    ppsci.utils.misc.set_random_seed(cfg.seed)
+    # initialize logger
+    logger.init_logger("ppsci", osp.join(cfg.output_dir, "eval.log"), "info")
+
+    mat = scipy.io.loadmat(cfg.data_file)
+    ag_data = mat["input_tf"]  # ag, ad, av
+    u_data = mat["target_X_tf"]
+    ut_data = mat["target_Xd_tf"]
+    utt_data = mat["target_Xdd_tf"]
+    ag_data = ag_data.reshape([ag_data.shape[0], ag_data.shape[1], 1])
+    u_data = u_data.reshape([u_data.shape[0], u_data.shape[1], 1])
+    ut_data = ut_data.reshape([ut_data.shape[0], ut_data.shape[1], 1])
+    utt_data = utt_data.reshape([utt_data.shape[0], utt_data.shape[1], 1])
+
+    t = mat["time"]
+    dt = t[0, 1] - t[0, 0]
+
+    ag_all = ag_data
+    u_all = u_data
+    u_t_all = ut_data
+    u_tt_all = utt_data
+
+    # finite difference
+    N = u_data.shape[1]
+    phi1 = np.concatenate(
+        [
+            np.array([-3 / 2, 2, -1 / 2]),
+            np.zeros([N - 3]),
+        ]
+    )
+    temp1 = np.concatenate([-1 / 2 * np.identity(N - 2), np.zeros([N - 2, 2])], axis=1)
+    temp2 = np.concatenate([np.zeros([N - 2, 2]), 1 / 2 * np.identity(N - 2)], axis=1)
+    phi2 = temp1 + temp2
+    phi3 = np.concatenate(
+        [
+            np.zeros([N - 3]),
+            np.array([1 / 2, -2, 3 / 2]),
+        ]
+    )
+    phi_t0 = (
+        1
+        / dt
+        * np.concatenate(
+            [
+                np.reshape(phi1, [1, phi1.shape[0]]),
+                phi2,
+                np.reshape(phi3, [1, phi3.shape[0]]),
+            ],
+            axis=0,
+        )
+    )
+    phi_t0 = np.reshape(phi_t0, [1, N, N])
+
+    ag_star = ag_all[0:10]
+    eta_star = u_all[0:10]
+    eta_t_star = u_t_all[0:10]
+    eta_tt_star = u_tt_all[0:10]
+    ag_c_star = ag_all[0:50]
+    lift_star = -ag_c_star
+
+    eta = eta_star
+    ag = ag_star
+    lift = lift_star
+    eta_t = eta_t_star
+    eta_tt = eta_tt_star
+    ag_c = ag_c_star
+    g = -eta_tt - ag
+    phi_t = np.repeat(phi_t0, ag_c_star.shape[0], axis=0)
+
+    model = ppsci.arch.DeepPhyLSTM(
+        cfg.MODEL.phylstm2_net.input_size,
+        eta.shape[2],
+        cfg.MODEL.phylstm2_net.hidden_size,
+        cfg.MODEL.phylstm2_net.model_type,
+    )
+    model.register_input_transform(functions.transform_in)
+    model.register_output_transform(functions.transform_out)
+
+    dataset_obj = functions.Dataset(eta, eta_t, g, ag, ag_c, lift, phi_t)
+    (
+        _,
+        _,
+        input_dict_val,
+        label_dict_val,
+    ) = dataset_obj.get(1)
+
+    sup_validator_pde = ppsci.validate.SupervisedValidator(
+        {
+            "dataset": {
+                "name": "NamedArrayDataset",
+                "input": input_dict_val,
+                "label": label_dict_val,
+            },
+        },
+        ppsci.loss.FunctionalLoss(functions.train_loss_func2),
+        {
+            "eta_pred": lambda out: out["eta_pred"],
+            "eta_dot_pred": lambda out: out["eta_dot_pred"],
+            "g_pred": lambda out: out["g_pred"],
+            "eta_t_pred_c": lambda out: out["eta_t_pred_c"],
+            "eta_dot_pred_c": lambda out: out["eta_dot_pred_c"],
+            "lift_pred_c": lambda out: out["lift_pred_c"],
+        },
+        metric={"metric": ppsci.metric.FunctionalMetric(functions.metric_expr)},
+        name="sup_valid",
+    )
+    validator_pde = {sup_validator_pde.name: sup_validator_pde}
+
+    # initialize solver
+    solver = ppsci.solver.Solver(
+        model,
+        output_dir=cfg.output_dir,
+        seed=cfg.seed,
+        validator=validator_pde,
+        checkpoint_path=cfg.EVAL.pretrained_model_path,
+        eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
+    )
+    # evaluate after finished training
+    solver.eval()
+
+
+@hydra.main(version_base=None, config_path="./conf", config_name="phylstm2.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()