[Fea] Support tensorboardX and add corresponding guidance (#812)

* support tensorboardX for viv as demo and add tensorboardX guide in
user_guide.md

* fix comma

Author: HydrogenSulfate

docs/zh/examples/viv.md

@@ -130,9 +130,9 @@ examples/fsi/viv.py:53:54
 
 接下来我们需要指定训练轮数和学习率，此处我们按实验经验，使用 10000 轮训练轮数，并每隔 10000 个epochs评估一次模型精度。
 
-``` yaml linenums="41"
+``` yaml linenums="42"
 --8<--
-examples/fsi/conf/viv.yaml:41:56
+examples/fsi/conf/viv.yaml:42:57
 --8<--
 ```
 

docs/zh/user_guide.md

@@ -274,7 +274,7 @@ pip install paddle2onnx
     [Paddle2ONNX] Start to parsing Paddle model...
     [Paddle2ONNX] Use opset_version = 13 for ONNX export.
     [Paddle2ONNX] PaddlePaddle model is exported as ONNX format now.
-    [2024/03/02 05:47:51] ppsci MESSAGE: ONNX model has been exported to: ./inference/aneurysm.onnx
+    ppsci MESSAGE: ONNX model has been exported to: ./inference/aneurysm.onnx
     ```
 
 ### 1.3 模型推理预测
@@ -410,6 +410,9 @@ PaddleScience 提供了多种推理配置组合，可通过命令行进行组合
     3. 运行 `aneurysm.py` 的推理功能，同时指定推理引擎为 TensorRT。
 
         ``` sh
+        # 运行前需设置指定GPU，否则可能无法启动 TensorRT
+        export CUDA_VISIBLE_DEVICES=0
+
         python aneurysm.py mode=infer \
             INFER.device=gpu \
             INFER.engine=tensorrt \
@@ -556,7 +559,47 @@ solver = ppsci.solver.Solver(
 solver.eval()
 ```
 
-### 1.7 使用 VisualDL 记录实验
+### 1.7 实验过程可视化
+
+#### 1.7.1 TensorBoardX
+
+[TensorBoardX](https://github.com/lanpa/tensorboardX) 是基于 TensorBoard 编写可视化分析工具，以丰富的图表呈现训练参数变化趋势、数据样本、模型结构、PR曲线、ROC曲线、高维数据分布等。帮助用户清晰直观地理解深度学习模型训练过程及模型结构，进而实现高效的模型调优。
+
+PaddleScience 支持使用 TensorBoardX 记录训练过程中的基础实验数据，包括 train/eval loss，eval metric，learning rate 等基本信息，可按如下步骤使用该功能。
+
+1. 安装 Tensorboard 和 TensorBoardX
+
+    ``` sh
+    pip install tensorboard tensorboardX
+    ```
+
+2. 在案例代码的 `Solver` 实例化时指定 `use_tbd=True`，然后再启动案例训练
+
+    ``` py hl_lines="3"
+    solver = ppsci.solver.Solver(
+        ...,
+        use_tbd=True,
+    )
+    ```
+
+3. 可视化记录数据
+
+    根据上述步骤，在训练时 TensorBoardX 会自动记录数据并保存到 `${solver.output_dir}/tensorboard` 目录下，具体所在路径在实例化 `Solver` 时，会自动打印在终端中，如下所示。
+
+    ``` log hl_lines="3" hl_lines="2"
+    ppsci MESSAGE: TensorboardX tool is enabled for logging, you can view it by running:
+    tensorboard --logdir outputs_VIV/2024-01-01/08-00-00/tensorboard
+    ```
+
+    !!! tip
+
+        也可以输入 `tensorboard --logdir ./outputs_VIV`，一次性在网页上展示 `outputs_VIV` 目录下所有训练记录，便于对比。
+
+    在终端里输入上述可视化命令，并用浏览器进入 TensorBoardX 给出的可视化地址，即可在浏览器内查看记录的数据，如下图所示。
+
+    ![tensorboardx_preview](https://paddle-org.bj.bcebos.com/paddlescience/docs/user_guide/tensorboardx_preview.JPG)
+
+#### 1.7.2 VisualDL
 
 [VisualDL](https://www.paddlepaddle.org.cn/paddle/visualdl) 是飞桨推出的可视化分析工具，以丰富的图表呈现训练参数变化趋势、数据样本、模型结构、PR曲线、ROC曲线、高维数据分布等。帮助用户清晰直观地理解深度学习模型训练过程及模型结构，进而实现高效的模型调优。
 
@@ -568,30 +611,31 @@ PaddleScience 支持使用 VisualDL 记录训练过程中的基础实验数据
     pip install -U visualdl
     ```
 
-2. 在案例代码的 `Solver` 实例化时指定 `use_visualdl=True`，然后再启动案例训练
+2. 在案例代码的 `Solver` 实例化时指定 `use_vdl=True`，然后再启动案例训练
 
     ``` py hl_lines="3"
     solver = ppsci.solver.Solver(
         ...,
-        use_visualdl=True,
+        use_vdl=True,
     )
     ```
 
 3. 可视化记录数据
 
-    根据上述步骤，在训练时 VisualDL 会自动记录数据并保存到 `${solver.output_dir}/vdl` 的目录中。`vdl` 所在路径在实例化 `Solver` 时，会自动打印在终端中，如下所示。
+    根据上述步骤，在训练时 VisualDL 会自动记录数据并保存到 `${solver.output_dir}/vdl` 目录下，具体所在路径在实例化 `Solver` 时，会自动打印在终端中，如下所示。
 
-    ``` log hl_lines="3"
+    ``` log hl_lines="4"
     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'.
+    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.8 使用 WandB 记录实验
+#### 1.7.3 WandB
 
 [WandB](https://wandb.ai/) 是一个第三方实验记录工具，能在记录实验数据的同时将数据上传到其用户的私人账户上，防止实验记录丢失。
 

examples/fsi/conf/viv.yaml

@@ -27,6 +27,7 @@ mode: train # running mode: train/eval
 seed: 42
 output_dir: ${hydra:run.dir}
 log_freq: 20
+use_tbd: false
 
 VIV_DATA_PATH: "./VIV_Training_Neta100.mat"
 

examples/fsi/viv.py

@@ -111,6 +111,7 @@ def train(cfg: DictConfig):
         lr_scheduler,
         cfg.TRAIN.epochs,
         cfg.TRAIN.iters_per_epoch,
+        use_tbd=cfg.use_tbd,
         save_freq=cfg.TRAIN.save_freq,
         log_freq=cfg.log_freq,
         eval_during_train=cfg.TRAIN.eval_during_train,

ppsci/solver/printer.py

@@ -103,6 +103,7 @@ def log_train_info(
         step=trainer.global_step,
         vdl_writer=trainer.vdl_writer,
         wandb_writer=trainer.wandb_writer,
+        tbd_writer=trainer.tbd_writer,
     )
 
 
@@ -145,4 +146,5 @@ def log_eval_info(
         step=trainer.global_step,
         vdl_writer=trainer.vdl_writer,
         wandb_writer=trainer.wandb_writer,
+        tbd_writer=trainer.tbd_writer,
     )

ppsci/solver/solver.py

@@ -70,6 +70,7 @@ class Solver:
         seed (int, optional): Random seed. Defaults to 42.
         use_vdl (Optional[bool]): Whether use VisualDL to log scalars. Defaults to False.
         use_wandb (Optional[bool]): Whether use wandb to log data. Defaults to False.
+        use_tbd (Optional[bool]): Whether use tensorboardX to log data. Defaults to False.
         wandb_config (Optional[Dict[str, str]]): Config dict of WandB. Defaults to None.
         device (Literal["cpu", "gpu", "xpu"], optional): Runtime device. Defaults to "gpu".
         equation (Optional[Dict[str, ppsci.equation.PDE]]): Equation dict. Defaults to None.
@@ -130,6 +131,7 @@ def __init__(
         seed: int = 42,
         use_vdl: bool = False,
         use_wandb: bool = False,
+        use_tbd: bool = False,
         wandb_config: Optional[Mapping] = None,
         device: Literal["cpu", "gpu", "xpu"] = "gpu",
         equation: Optional[Dict[str, ppsci.equation.PDE]] = None,
@@ -337,8 +339,8 @@ def dist_wrapper(model: nn.Layer) -> paddle.DataParallel:
                 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'."
+                "VisualDL is enabled for logging, you can view it by "
+                f"running:\nvisualdl --logdir {self.vdl_writer._logdir} --port 8080"
             )
 
         # set WandB tool
@@ -354,6 +356,25 @@ def dist_wrapper(model: nn.Layer) -> paddle.DataParallel:
                 if is_master:
                     self.wandb_writer = wandb.init(**wandb_config)
 
+        # set TensorBoardX tool
+        self.tbd_writer = None
+        if use_tbd:
+            try:
+                import tensorboardX
+            except ModuleNotFoundError:
+                raise ModuleNotFoundError(
+                    "Please install 'tensorboardX' with `pip install tensorboardX` first."
+                )
+            with misc.RankZeroOnly(self.rank) as is_master:
+                if is_master:
+                    self.tbd_writer = tensorboardX.SummaryWriter(
+                        osp.join(output_dir, "tensorboard")
+                    )
+            logger.message(
+                "TensorboardX is enabled for logging, you can view it by "
+                f"running:\ntensorboard --logdir {self.tbd_writer.logdir}"
+            )
+
         self.global_step = 0
 
         # log paddlepaddle's version
@@ -462,6 +483,7 @@ def train(self) -> None:
                         epoch_id,
                         self.vdl_writer,
                         self.wandb_writer,
+                        self.tbd_writer,
                     )
 
                 # visualize after evaluation

ppsci/utils/logger.py

@@ -31,6 +31,7 @@
 if TYPE_CHECKING:
     import visualdl  # isort:skip
     import wandb  # isort:skip
+    import tensorboardX as tbd
 
 _logger: logging.Logger = None
 
@@ -200,6 +201,7 @@ def scalar(
     step: int,
     vdl_writer: Optional["visualdl.LogWriter"] = None,
     wandb_writer: Optional["wandb.run"] = None,
+    tbd_writer: Optional["tbd.SummaryWriter"] = None,
 ):
     """This function will add scalar data to VisualDL or WandB for plotting curve(s).
 
@@ -210,14 +212,22 @@ def scalar(
         wandb_writer (wandb.run): Run object of WandB to record metrics. Defaults to None.
     """
     if vdl_writer is not None:
-        for name, value in metric_dict.items():
-            vdl_writer.add_scalar(name, value, step)
+        with misc.RankZeroOnly() as is_master:
+            if is_master:
+                for name, value in metric_dict.items():
+                    vdl_writer.add_scalar(name, value, step)
 
     if wandb_writer is not None:
         with misc.RankZeroOnly() as is_master:
             if is_master:
                 wandb_writer.log({"step": step, **metric_dict})
 
+    if tbd_writer is not None:
+        with misc.RankZeroOnly() as is_master:
+            if is_master:
+                for name, value in metric_dict.items():
+                    tbd_writer.add_scalar(name, value, global_step=step)
+
 
 def advertise():
     """

ppsci/utils/symbolic.py

@@ -107,19 +107,6 @@
 }
 
 
-def _numerator_of_derivative(expr: sp.Basic) -> sp.Basic:
-    if not isinstance(expr, sp.Derivative):
-        raise TypeError(
-            f"expr({expr}) should be of type sp.Derivative, but got {type(expr)}"
-        )
-    if len(expr.args) <= 2:
-        if expr.args[1][1] == 1:
-            return expr.args[0]
-        return sp.Derivative(expr.args[0], (expr.args[1][0], expr.args[1][1] - 1))
-    else:
-        return sp.Derivative(*expr.args[:-1])
-
-
 def _cvt_to_key(expr: sp.Basic) -> str:
     """Convert sympy expression to a string key, mainly as retrieval key in dict.
 
@@ -585,7 +572,7 @@ def _visualize_graph(nodes: List[sp.Basic], graph_filename: str):
     }
     naming_counter = {k: 0 for k in SYMPY_BUILTIN_NAME}
 
-    def get_operator_name(node):
+    def get_operator_name(node: sp.Function):
         ret = f"{SYMPY_BUILTIN_NAME[node.func]}_{naming_counter[node.func]}"
         naming_counter[node.func] += 1
         return ret
@@ -601,8 +588,8 @@ def add_edge(u: str, v: str, u_color: str = C_DATA, v_color: str = C_DATA):
         Args:
             u (str): Name of begin node u.
             v (str): Name of end node v.
-            u_color (str, optional): _description_. Defaults to C_DATA.
-            v_color (str, optional): _description_. Defaults to C_DATA.
+            u_color (str, optional): Color of node u. Defaults to '#feb64d'.
+            v_color (str, optional): Color of node v. Defaults to '#feb64d'.
         """
         graph.add_node(u, style="filled", shape="ellipse", color=u_color)
         graph.add_node(v, style="filled", shape="ellipse", color=v_color)