[Fix&Doc] Fix vdl (#671)

* fix wrong argument order for vdl and add note when vdl is enabled

* update user_guide for VDL tool and refine some code

* add accurate condition for saving *.eqn file, refine code in init_logger

Author: HydrogenSulfate

docs/requirements.txt

@@ -6,7 +6,6 @@ mkdocs-git-revision-date-localized-plugin
 mkdocs-glightbox
 mkdocs-include-markdown-plugin
 mkdocs-material
-mkdocs-material
 mkdocs-material-extensions
 mkdocs-video
 mkdocstrings

docs/zh/user_guide.md

@@ -169,7 +169,7 @@ PaddleScience/examples/bracket/outputs_bracket/
 
 1. 加载 `*.pdparams` 文件内的参数到模型中
 
-    ``` py
+    ``` py hl_lines="9 10"
     import ppsci
     import numpy as np
 
@@ -186,7 +186,7 @@ PaddleScience/examples/bracket/outputs_bracket/
 
 2. 准备好用于预测的输入数据，并以字典 `dict` 的方式传递给 `solver.predict`。
 
-    ``` py
+    ``` py hl_lines="12 13 14 15 16"
     N = 100 # 假设要预测100个样本的结果
     x = np.random.randn(N, 1) # 准备 字段
     y = np.random.randn(N, 1)
@@ -224,7 +224,7 @@ PaddleScience/examples/bracket/outputs_bracket/
 
 因此我们只需要在 `Solver` 时指定 `checkpoint_path` 参数为 `latest.*` 的所在路径，即可自动载入上述的几个文件，并从 `latest` 中记录的 epoch 开始继续训练。
 
-``` py hl_lines="9"
+``` py hl_lines="7"
 import ppsci
 
 ...
@@ -241,19 +241,35 @@ solver = ppsci.solver.Solver(
 
 ### 1.4 迁移学习
 
-迁移学习是一种广泛使用、低成本提高模型精度的训练方式。在 PaddleScience 中，只需在 `model` 实例化完毕之后，手动为其载入预训练模型权重，即可进行迁移学习。
+迁移学习是一种广泛使用、低成本提高模型精度的训练方式。在 PaddleScience 中，可以通过在 `model` 实例化完毕之后，手动为其载入预训练模型权重；也可以在 `Solver` 实例化时指定 `pretrained_model_path` 自动载入预训练模型权重，两种方式都可以进行迁移学习。
 
-``` py hl_lines="9"
-import ppsci
-import ppsci.utils
-from ppsci.utils import save_load
+=== "手动载入预训练模型"
 
-...
-...
+    ``` py hl_lines="8"
+    import ppsci
+    from ppsci.utils import save_load
 
-model = ...
-save_load.load_pretrain(model, "/path/to/pretrain")
-```
+    ...
+    ...
+
+    model = ...
+    save_load.load_pretrain(model, "/path/to/pretrain")
+    ```
+
+=== "指定 `pretrained_model_path` 自动载入预训练模型"
+
+    ``` py hl_lines="9"
+    import ppsci
+
+    ...
+    ...
+
+    model = ...
+    solver = ppsci.solver.Solver(
+        ...,
+        pretrained_model_path="/path/to/pretrain",
+    )
+    ```
 
 !!! info "迁移学习建议"
 
@@ -278,7 +294,42 @@ solver = ppsci.solver.Solver(
 solver.eval()
 ```
 
-### 1.6 使用 WandB 记录实验
+### 1.6 使用 VisualDL 记录实验
+
+[VisualDL](https://www.paddlepaddle.org.cn/paddle/visualdl) 是飞桨推出的可视化分析工具，以丰富的图表呈现训练参数变化趋势、数据样本、模型结构、PR曲线、ROC曲线、高维数据分布等。帮助用户清晰直观地理解深度学习模型训练过程及模型结构，进而实现高效的模型调优。
+
+PaddleScience 支持使用 VisualDL 记录训练过程中的基础实验数据，包括 train/eval loss，eval metric，learning rate 等基本信息，可按如下步骤使用该功能。
+
+1. 安装 VisualDL
+
+    ``` sh
+    pip install -U visualdl
+    ```
+
+2. 在案例代码的 `Solver` 实例化时指定 `use_visualdl=True`，然后再启动案例训练
+
+    ``` py hl_lines="3"
+    solver = ppsci.solver.Solver(
+        ...,
+        use_visualdl=True,
+    )
+    ```
+
+3. 可视化记录数据
+
+    根据上述步骤，在训练时 VisualDL 会自动记录数据并保存到 `${solver.output_dir}/vdl` 的目录中。`vdl` 所在路径在实例化 `Solver` 时，会自动打印在终端中，如下所示。
+
+    ``` log hl_lines="3"
+    Please NOTE: device: 0, GPU Compute Capability: 7.0, Driver API Version: 11.8, Runtime API Version: 11.6
+    device: 0, cuDNN Version: 8.4.
+    ppsci INFO: VisualDL tool enabled for logging, you can view it by running: 'visualdl --logdir outputs_darcy2d/2023-10-08/10-00-00/TRAIN.epochs=400/vdl --port 8080'.
+    ```
+
+    在终端里输入上述可视化命令，并用浏览器进入 VisualDL 给出的可视化地址，即可在浏览器内查看记录的数据，如下图所示。
+
+    ![visualdl_record](https://paddle-org.bj.bcebos.com/paddlescience/docs/user_guide/VisualDL_preview.png)
+
+### 1.7 使用 WandB 记录实验
 
 [WandB](https://wandb.ai/) 是一个第三方实验记录工具，能在记录实验数据的同时将数据上传到其用户的私人账户上，防止实验记录丢失。
 
@@ -331,7 +382,7 @@ PaddleScience 支持使用 WandB 记录基本的实验数据，包括 train/eval
 
 接下来以 `examples/pipe/poiseuille_flow.py` 为例，介绍如何正确使用 PaddleScience 的数据并行功能。分布式训练细节可以参考：[Paddle-使用指南-分布式训练-快速开始-数据并行](https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/guides/06_distributed_training/cluster_quick_start_collective_cn.html)。
 
-1. 在 constraint 实例化完毕后，将 `ITERS_PER_EPOCH` 重新赋值为经过自动多卡数据切分后的 `dataloader` 的长度（一般情况下其长度等于单卡 dataloader 的长度除以卡数，向上取整），如代码中黄色高亮行所示。
+1. 在 constraint 实例化完毕后，将 `ITERS_PER_EPOCH` 重新赋值为经过自动多卡数据切分后的 `dataloader` 的长度（一般情况下其长度等于单卡 dataloader 的长度除以卡数，向上取整），如代码中高亮行所示。
 
     ``` py linenums="146" title="examples/pipe/poiseuille_flow.py" hl_lines="22"
     ITERS_PER_EPOCH = int((N_x * N_y * N_p) / BATCH_SIZE)
@@ -379,7 +430,7 @@ TODO -->
 
 接下来介绍如何正确使用 PaddleScience 的自动混合精度功能。自动混合精度的原理可以参考：[Paddle-使用指南-性能调优-自动混合精度训练（AMP）](https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/guides/performance_improving/amp_cn.html#amp)。
 
-实例化 `Solver` 时加上 2 个参数: `use_amp=True`, `amp_level="O1"`(或`amp_level="O2"`)。如代码中黄色高亮行所示，通过指定 `use_amp=True`，开启自动混合精度功能，接着再设置 `amp_level="O1"`，指定混合精度所用的模式，`O1` 为自动混合精度，`O2` 为更激进的纯 fp16 训练模式，一般推荐使用 `O1`。
+实例化 `Solver` 时加上 2 个参数: `use_amp=True`, `amp_level="O1"`(或`amp_level="O2"`)。如代码中高亮行所示，通过指定 `use_amp=True`，开启自动混合精度功能，接着再设置 `amp_level="O1"`，指定混合精度所用的模式，`O1` 为自动混合精度，`O2` 为更激进的纯 fp16 训练模式，一般推荐使用 `O1`。
 
 ``` py hl_lines="5 6"
 # initialize solver
@@ -395,7 +446,7 @@ solver = ppsci.solver.Solver(
 
 接下来介绍如何正确使用 PaddleScience 的梯度累加功能。梯度累加的原理可以参考：[Paddle-使用指南-性能调优-自动混合精度训练（AMP）-动态图下使用梯度累加](https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/guides/performance_improving/amp_cn.html#dongtaituxiashiyongtiduleijia)。
 
-实例化 `Solver` 时指定 `update_freq` 参数为大于 1 的正整数即可。如代码中黄色高亮行所示，`update_freq` 可以设置为 2 或者更大的整数，推荐使用 2、4、8，此时对于训练任务来说，全局 `batch size` 等价于 `update_freq * batch size`。梯度累加方法在大多数场景中能够让间接地扩大每个 batch 内的样本数量，从而让每个 batch 分布更接近真实数据分布，提升训练任务的性能。
+实例化 `Solver` 时指定 `update_freq` 参数为大于 1 的正整数即可。如代码中高亮行所示，`update_freq` 可以设置为 2 或者更大的整数，推荐使用 2、4、8，此时对于训练任务来说，全局 `batch size` 等价于 `update_freq * batch size`。梯度累加方法在大多数场景中能够让间接地扩大每个 batch 内的样本数量，从而让每个 batch 分布更接近真实数据分布，提升训练任务的性能。
 
 ``` py hl_lines="5"
 # initialize solver

examples/amgnet/amgnet_airfoil.py

@@ -136,6 +136,7 @@ def train(cfg: DictConfig):
     solver.train()
 
     # visualize prediction
+    logger.message("Now visualizing prediction, please wait...")
     with solver.no_grad_context_manager(True):
         for index, (input_, label, _) in enumerate(rmse_validator.data_loader):
             truefield = label["label"].y

examples/amgnet/amgnet_cylinder.py

@@ -136,6 +136,7 @@ def train(cfg: DictConfig):
     solver.train()
 
     # visualize prediction
+    logger.message("Now visualizing prediction, please wait...")
     with solver.no_grad_context_manager(True):
         for index, (input_, label, _) in enumerate(rmse_validator.data_loader):
             truefield = label["label"].y

ppsci/solver/solver.py

@@ -300,6 +300,10 @@ def __init__(
             with misc.RankZeroOnly(self.rank) as is_master:
                 if is_master:
                     self.vdl_writer = vdl.LogWriter(osp.join(output_dir, "vdl"))
+            logger.info(
+                "VisualDL tool is enabled for logging, you can view it by "
+                f"running: 'visualdl --logdir {self.vdl_writer._logdir} --port 8080'."
+            )
 
         # set WandB tool
         self.wandb_writer = None
@@ -413,7 +417,10 @@ def train(self):
                 )
                 for metric_dict in metric_dict_group.values():
                     logger.scaler(
-                        metric_dict, epoch_id, self.vdl_writer, self.wandb_writer
+                        {f"eval/{k}": v for k, v in metric_dict.items()},
+                        epoch_id,
+                        self.vdl_writer,
+                        self.wandb_writer,
                     )
 
                 # visualize after evaluation

ppsci/utils/logger.py

@@ -100,9 +100,9 @@ def init_logger(
     stream_handler._name = "stream_handler"
     _logger.addHandler(stream_handler)
 
-    # add file_handler, output to log_file(if specified)
+    # add file_handler, output to log_file(if specified), only for rank 0 device
     if log_file is not None and dist.get_rank() == 0:
-        log_file_folder = os.path.split(log_file)[0]
+        log_file_folder = os.path.dirname(log_file)
         os.makedirs(log_file_folder, exist_ok=True)
         file_formatter = logging.Formatter(
             "[%(asctime)s] %(name)s %(levelname)s: %(message)s",
@@ -122,7 +122,18 @@ def init_logger(
 
 
 def set_log_level(log_level: int):
-    """Set logger level, only msg of level >= `log_level` will be printed.
+    """Set logger level, only message of level >= `log_level` will be printed.
+
+    Built-in log level are below:
+
+    CRITICAL = 50,
+    FATAL = 50,
+    ERROR = 40,
+    WARNING = 30,
+    WARN = 30,
+    INFO = 20,
+    DEBUG = 10,
+    NOTSET = 0.
 
     Args:
         log_level (int): Log level.
@@ -135,7 +146,7 @@ def set_log_level(log_level: int):
 
 def ensure_logger(log_func: Callable) -> Callable:
     """
-    Automatically initialize `logger` by default arguments
+    A decorator which automatically initialize `logger` by default arguments
     when init_logger() is not called manually.
     """
 
@@ -201,7 +212,7 @@ def scaler(
     """
     if vdl_writer is not None:
         for name, value in metric_dict.items():
-            vdl_writer.add_scalar(name, step, value)
+            vdl_writer.add_scalar(name, value, step)
 
     if wandb_writer is not None:
         with misc.RankZeroOnly() as is_master:

ppsci/utils/save_load.py

@@ -186,9 +186,13 @@ def save_checkpoint(
     if grad_scaler is not None:
         paddle.save(grad_scaler.state_dict(), f"{ckpt_path}.pdscaler")
     if equation is not None:
-        paddle.save(
-            {key: eq.state_dict() for key, eq in equation.items()},
-            f"{ckpt_path}.pdeqn",
+        num_learnable_params = sum(
+            [len(eq.learnable_parameters) for eq in equation.values()]
         )
+        if num_learnable_params > 0:
+            paddle.save(
+                {key: eq.state_dict() for key, eq in equation.items()},
+                f"{ckpt_path}.pdeqn",
+            )
 
     logger.message(f"Finish saving checkpoint to {ckpt_path}")