Merge branch 'develop' into support_nvtx

Author: HydrogenSulfate

docs/zh/examples/laplace2d.md

@@ -60,9 +60,9 @@ $$
 
 上式中 $f$ 即为 MLP 模型本身，用 PaddleScience 代码表示如下
 
-``` py linenums="32"
+``` py linenums="23"
 --8<--
-examples/laplace/laplace2d.py:32:33
+examples/laplace/laplace2d.py:23:24
 --8<--
 ```
 
@@ -74,9 +74,9 @@ examples/laplace/laplace2d.py:32:33
 
 由于 2D-Laplace 使用的是 Laplace 方程的2维形式，因此可以直接使用 PaddleScience 内置的 `Laplace`，指定该类的参数 `dim` 为2。
 
-``` py linenums="35"
+``` py linenums="26"
 --8<--
-examples/laplace/laplace2d.py:35:36
+examples/laplace/laplace2d.py:26:27
 --8<--
 ```
 
@@ -85,9 +85,9 @@ examples/laplace/laplace2d.py:35:36
 本文中 2D Laplace 问题作用在以 (0.0, 0.0),  (1.0, 1.0) 为对角线的二维矩形区域，
 因此可以直接使用 PaddleScience 内置的空间几何 `Rectangle` 作为计算域。
 
-``` py linenums="38"
+``` py linenums="29"
 --8<--
-examples/laplace/laplace2d.py:38:43
+examples/laplace/laplace2d.py:29:34
 --8<--
 ```
 
@@ -97,19 +97,19 @@ examples/laplace/laplace2d.py:38:43
 
 在定义约束之前，需要给每一种约束指定采样点个数，表示每一种约束在其对应计算域内采样数据的数量，以及通用的采样配置。
 
-``` yaml linenums="26"
+``` yaml linenums="30"
 --8<--
-examples/laplace/conf/laplace2d.yaml:26:27
+examples/laplace/conf/laplace2d.yaml:30:31
 --8<--
 ```
 
 #### 3.4.1 内部点约束
 
 以作用在内部点上的 `InteriorConstraint` 为例，代码如下：
 
-``` py linenums="60"
+``` py linenums="50"
 --8<--
-examples/laplace/laplace2d.py:60:68
+examples/laplace/laplace2d.py:50:59
 --8<--
 ```
 
@@ -131,19 +131,19 @@ examples/laplace/laplace2d.py:60:68
 
 同理，我们还需要构建矩形的四个边界的约束。但与构建 `InteriorConstraint` 约束不同的是，由于作用区域是边界，因此我们使用 `BoundaryConstraint` 类，代码如下：
 
-``` py linenums="69"
+``` py linenums="60"
 --8<--
-examples/laplace/laplace2d.py:69:76
+examples/laplace/laplace2d.py:60:72
 --8<--
 ```
 
 `BoundaryConstraint` 类第一个参数表示我们直接对网络模型的输出结果 `out["u"]` 作为程序运行时的约束对象；
 
 第二个参数是指我们约束对象的真值如何获得，这里我们直接通过其解析解进行计算，定义解析解的代码如下：
 
-``` py linenums="45"
+``` py linenums="36"
 --8<--
-examples/laplace/laplace2d.py:45:49
+examples/laplace/laplace2d.py:36:40
 --8<--
 ```
 
@@ -153,29 +153,29 @@ examples/laplace/laplace2d.py:45:49
 
 接下来我们需要在配置文件中指定训练轮数，此处我们按实验经验，使用两万轮训练轮数，评估间隔为两百轮。
 
-``` yaml linenums="41"
+``` yaml linenums="45"
 --8<--
-examples/laplace/conf/laplace2d.yaml:41:46
+examples/laplace/conf/laplace2d.yaml:45:50
 --8<--
 ```
 
 ### 3.6 优化器构建
 
 训练过程会调用优化器来更新模型参数，此处选择较为常用的 `Adam` 优化器。
 
-``` py linenums="83"
+``` py linenums="74"
 --8<--
-examples/laplace/laplace2d.py:83:84
+examples/laplace/laplace2d.py:74:75
 --8<--
 ```
 
 ### 3.7 评估器构建
 
 在训练过程中通常会按一定轮数间隔，用验证集（测试集）评估当前模型的训练情况，因此使用 `ppsci.validate.GeometryValidator` 构建评估器。
 
-``` py linenums="86"
+``` py linenums="77"
 --8<--
-examples/laplace/laplace2d.py:86:100
+examples/laplace/laplace2d.py:77:92
 --8<--
 ```
 
@@ -185,19 +185,19 @@ examples/laplace/laplace2d.py:86:100
 
 本文中的输出数据是一个区域内的二维点集，因此我们只需要将评估的输出数据保存成 **vtu格式** 文件，最后用可视化软件打开查看即可。代码如下：
 
-``` py linenums="103"
+``` py linenums="94"
 --8<--
-examples/laplace/laplace2d.py:103:112
+examples/laplace/laplace2d.py:94:103
 --8<--
 ```
 
 ### 3.9 模型训练、评估与可视化
 
 完成上述设置之后，只需要将上述实例化的对象按顺序传递给 `ppsci.solver.Solver`，然后启动训练、评估、可视化。
 
-``` py linenums="114"
+``` py linenums="105"
 --8<--
-examples/laplace/laplace2d.py:114:134
+examples/laplace/laplace2d.py:105:125
 --8<--
 ```
 

examples/laplace/conf/laplace2d.yaml

@@ -14,14 +14,17 @@ hydra:
           - mode
           - output_dir
           - log_freq
+  callbacks:
+    init_callback:
+      _target_: ppsci.utils.callbacks.InitCallback
   sweep:
     # output directory for multirun
     dir: ${hydra.run.dir}
     subdir: ./
 
 # general settings
 mode: train # running mode: train/eval
-seed: 42
+seed: 2024
 log_freq: 20
 output_dir: ${hydra:run.dir}
 NPOINT_INTERIOR: 9801

examples/laplace/laplace2d.py

@@ -12,23 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from os import path as osp
-
 import hydra
 import numpy as np
 from omegaconf import DictConfig
 
 import ppsci
-from ppsci.utils import logger
 
 
 def train(cfg: DictConfig):
-    # set random seed for reproducibility
-    ppsci.utils.misc.set_random_seed(cfg.seed)
-
-    # # set output directory
-    logger.init_logger("ppsci", osp.join(cfg.output_dir, "train.log"), "info")
-
     # set model
     model = ppsci.arch.MLP(**cfg.MODEL)
 
@@ -135,11 +126,6 @@ def u_solution_func(out):
 
 
 def evaluate(cfg: DictConfig):
-    # set random seed for reproducibility
-    ppsci.utils.misc.set_random_seed(cfg.seed)
-    # set output directory
-    logger.init_logger("ppsci", osp.join(cfg.output_dir, "eval.log"), "info")
-
     # set model
     model = ppsci.arch.MLP(**cfg.MODEL)
 

ppsci/data/dataset/radar_dataset.py

@@ -123,7 +123,9 @@ def __getitem__(self, index):
         mask[data < 0] = 0
         data[data < 0] = 0
         data = np.clip(data, 0, 128)
-        vid = np.zeros((self.length, self.img_height, self.img_width, 2))
+        vid = np.zeros(
+            (self.length, self.img_height, self.img_width, 2), dtype=self.data_type
+        )
         vid[..., 0] = data
         vid[..., 1] = mask
 

ppsci/equation/pde/base.py

@@ -83,19 +83,71 @@ def add_equation(self, name: str, equation: Callable):
         self.equations.update({name: equation})
 
     def parameters(self) -> List[paddle.Tensor]:
-        """Return parameters contained in PDE.
+        """Return learnable parameters contained in PDE.
+
+        Args:
+            None
 
         Returns:
-            List[Tensor]: A list of parameters.
+            List[Tensor]: A list of learnable parameters.
+
+        Examples:
+            >>> import ppsci
+            >>> pde = ppsci.equation.Vibration(2, -4, 0)
+            >>> print(pde.parameters())
+            [Parameter containing:
+            Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=False,
+                   -4.), Parameter containing:
+            Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=False,
+                   0.)]
         """
         return self.learnable_parameters.parameters()
 
     def state_dict(self) -> Dict[str, paddle.Tensor]:
-        """Return named parameters in dict."""
+        """Return named learnable parameters in dict.
+
+        Args:
+            None
+
+        Returns:
+            Dict[str, Tensor]: A dict of states(str) and learnable parameters(Tensor).
+
+        Examples:
+            >>> import ppsci
+            >>> pde = ppsci.equation.Vibration(2, -4, 0)
+            >>> print(pde.state_dict())
+            OrderedDict([('0', Parameter containing:
+            Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=False,
+                   -4.)), ('1', Parameter containing:
+            Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=False,
+                   0.))])
+        """
+
         return self.learnable_parameters.state_dict()
 
-    def set_state_dict(self, state_dict):
-        """Set state dict from dict."""
+    def set_state_dict(self, state_dict: Dict[str, paddle.Tensor]):
+        """Set state dict from dict.
+
+        Args:
+            state_dict (Dict[str, paddle.Tensor]): The state dict to be set.
+
+        Returns:
+            None
+
+        Examples:
+            >>> import paddle
+            >>> import ppsci
+            >>> paddle.set_default_dtype("float64")
+            >>> pde = ppsci.equation.Vibration(2, -4, 0)
+            >>> state = pde.state_dict()
+            >>> state['0'] = paddle.to_tensor(-3.1)
+            >>> pde.set_state_dict(state)
+            >>> print(state)
+            OrderedDict([('0', Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=True,
+                   -3.10000000)), ('1', Parameter containing:
+            Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=False,
+                   0.))])
+        """
         self.learnable_parameters.set_state_dict(state_dict)
 
     def __str__(self):