[Fea&Fix] Support register dataset and fix several bugs (#1207)

HydrogenSulfate · web-flow · commit efa43c007d56 · 2025-08-22T17:27:45.000+08:00
* support register custom dataset to ppsci.data.dataset

* fix typing for cfg

* fix typing for cfg

* fix several kernel not supported in xpu

* fix drivaernet datasets
diff --git a/competition/IJCAI_2024_CAR b/competition/IJCAI_2024_CAR
@@ -1 +1 @@
-Subproject commit 4e12f693dbe0f84fc37c03f716a7353280d52a83
+Subproject commit 24adee632060693a9fdadb31c9880c170f94921d
diff --git a/ppsci/arch/__init__.py b/ppsci/arch/__init__.py
@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import copy
+from typing import TYPE_CHECKING
 
 from ppsci.arch.afno import AFNONet  # isort:skip
 from ppsci.arch.afno import PrecipNet  # isort:skip
@@ -64,6 +65,10 @@
 from ppsci.arch.ifm_mlp import IFMMLP  # isort:skip
 from ppsci.arch.stafnet import STAFNet  # isort:skip
 
+if TYPE_CHECKING:
+    from omegaconf import DictConfig
+
+
 __all__ = [
     "MoFlowNet",
     "MoFlowProp",
@@ -118,7 +123,7 @@
 ]
 
 
-def build_model(cfg):
+def build_model(cfg: DictConfig):
     """Build model
 
     Args:
diff --git a/ppsci/arch/afno.py b/ppsci/arch/afno.py
@@ -203,6 +203,19 @@ def __init__(
             default_initializer=nn.initializer.Normal(std=self.scale),
         )
 
+        if paddle.device.is_compiled_with_xpu():
+            # softshrink is not supported in xpu, so we use a custom implementation to avoid fallback to cpu
+            def softshrink(x: paddle.Tensor, threshold: float = 0.5) -> paddle.Tensor:
+                return paddle.where(
+                    x > threshold,
+                    x - threshold,
+                    paddle.where(x < -threshold, x + threshold, paddle.zeros_like(x)),
+                )
+
+            self.softshrink = softshrink
+        else:
+            self.softshrink = F.softshrink
+
     def forward(self, x):
         bias = x
 
@@ -285,7 +298,7 @@ def forward(self, x):
         )
 
         x = paddle.stack([o2_real, o2_imag], axis=-1)
-        x = F.softshrink(x, threshold=self.sparsity_threshold)
+        x = self.softshrink(x, threshold=self.sparsity_threshold)
         x = paddle.as_complex(x)
         x = x.reshape((B, H, W // 2 + 1, C))
         x = paddle.fft.irfft2(x, s=(H, W), axes=(1, 2), norm="ortho")
diff --git a/ppsci/constraint/__init__.py b/ppsci/constraint/__init__.py
@@ -42,7 +42,7 @@ def build_constraint(cfg, equation_dict, geom_dict):
     """Build constraint(s).
 
     Args:
-        cfg (List[DictConfig]): Constraint config list.
+        cfg (DictConfig): Constraint config list.
         equation_dict (Dct[str, Equation]): Equation(s) in dict.
         geom_dict (Dct[str, Geometry]): Geometry(ies) in dict.
 
diff --git a/ppsci/data/__init__.py b/ppsci/data/__init__.py
@@ -26,6 +26,7 @@
 from ppsci.data import dataloader
 from ppsci.data import dataset
 from ppsci.data import process
+from ppsci.data.dataset import register_to_dataset
 from ppsci.data.process import batch_transform
 from ppsci.data.process import transform
 from ppsci.utils import logger
@@ -37,6 +38,7 @@
     "build_dataloader",
     "transform",
     "batch_transform",
+    "register_to_dataset",
 ]
 
 
diff --git a/ppsci/data/dataset/__init__.py b/ppsci/data/dataset/__init__.py
@@ -12,9 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
+import sys
 from typing import TYPE_CHECKING
 
+from paddle import io
+
 from ppsci.data.dataset.airfoil_dataset import MeshAirfoilDataset
 from ppsci.data.dataset.array_dataset import ChipHeatDataset
 from ppsci.data.dataset.array_dataset import ContinuousNamedArrayDataset
@@ -56,7 +61,7 @@
 from ppsci.utils import logger
 
 if TYPE_CHECKING:
-    from paddle import io
+    from omegaconf import DictConfig
 
 __all__ = [
     "IterableNamedArrayDataset",
@@ -97,14 +102,15 @@
     "IFMMoeDataset",
     "STAFNetDataset",
     "TMTDataset",
+    "register_to_dataset",
 ]
 
 
-def build_dataset(cfg) -> "io.Dataset":
+def build_dataset(cfg: DictConfig) -> "io.Dataset":
     """Build dataset
 
     Args:
-        cfg (List[DictConfig]): Dataset config list.
+        cfg (DictConfig): Dataset config list.
 
     Returns:
         Dict[str, io.Dataset]: dataset.
@@ -115,8 +121,37 @@ def build_dataset(cfg) -> "io.Dataset":
     if "transforms" in cfg:
         cfg["transforms"] = transform.build_transforms(cfg.pop("transforms"))
 
-    dataset = eval(dataset_cls)(**cfg)
+    try:
+        dataset = eval(dataset_cls)(**cfg)
+    except NameError:
+        import textwrap
+
+        logger.error(
+            f"name {dataset_cls} is not defined, maybe you should register your dataset class first as below:\n"
+            + textwrap.indent(
+                "\n"
+                "import paddle\n"
+                "from ppsci.data import register_to_dataset\n"
+                "\n"
+                "@register_to_dataset\n"
+                "class MyDataset(paddle.io.Dataset):\n"
+                "    pass\n"
+                "\n",
+                prefix=" " * 4,
+            )
+        )
+        raise
 
     logger.debug(str(dataset))
 
     return dataset
+
+
+def register_to_dataset(cls: type):
+    from ppsci.utils.registry import register_cls_to_module
+
+    if not issubclass(cls, io.Dataset):
+        logger.warning(
+            f"The registered class '{cls.__name__}' should be inherited from `paddle.io.Dataset`"
+        )
+    register_cls_to_module(sys.modules[__name__], cls)
diff --git a/ppsci/data/dataset/drivaernet_dataset.py b/ppsci/data/dataset/drivaernet_dataset.py
@@ -245,7 +245,7 @@ def _load_point_cloud(self, design_id: str) -> Optional[paddle.Tensor]:
         load_path = os.path.join(self.root_dir, f"{design_id}.paddle_tensor")
         if os.path.exists(load_path) and os.path.getsize(load_path) > 0:
             try:
-                vertices = paddle.load(path=str(load_path))
+                vertices: paddle.Tensor = paddle.load(path=str(load_path))
                 num_vertices = vertices.shape[0]
 
                 if num_vertices > self.num_points:
@@ -255,6 +255,8 @@ def _load_point_cloud(self, design_id: str) -> Optional[paddle.Tensor]:
                     vertices = vertices.numpy()[indices]
                     vertices = paddle.to_tensor(vertices)
 
+                vertices = self._sample_or_pad_vertices(vertices, self.num_points)
+
                 return vertices
             except (EOFError, RuntimeError, ValueError) as e:
                 raise Exception(
diff --git a/ppsci/data/dataset/drivaernetplusplus_dataset.py b/ppsci/data/dataset/drivaernetplusplus_dataset.py
@@ -246,7 +246,7 @@ def _load_point_cloud(self, design_id: str):
         load_path = os.path.join(self.root_dir, f"{design_id}.paddle_tensor")
         if os.path.exists(load_path) and os.path.getsize(load_path) > 0:
             try:
-                vertices = paddle.load(path=str(load_path))
+                vertices: paddle.Tensor = paddle.load(path=str(load_path))
             except (EOFError, RuntimeError, ValueError) as e:
                 raise Exception(
                     f"Error loading point cloud from {load_path}: {e}"
@@ -256,6 +256,9 @@ def _load_point_cloud(self, design_id: str):
             if num_vertices > self.num_points:
                 indices = np.random.choice(num_vertices, self.num_points, replace=False)
                 vertices = vertices.numpy()[indices]
+                vertices = paddle.to_tensor(vertices)
+
+            vertices = self._sample_or_pad_vertices(vertices, self.num_points)
 
             return vertices
 
diff --git a/ppsci/data/process/batch_transform/__init__.py b/ppsci/data/process/batch_transform/__init__.py
@@ -12,10 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
 import numbers
 from collections.abc import Mapping
 from collections.abc import Sequence
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
 from typing import List
@@ -27,6 +30,10 @@
 from ppsci.data.process import transform
 from ppsci.data.process.batch_transform.preprocess import FunctionalBatchTransform
 
+if TYPE_CHECKING:
+    from omegaconf import DictConfig
+
+
 try:
     import pgl
 except ModuleNotFoundError:
@@ -104,7 +111,7 @@ def default_collate_fn(batch: List[Any]) -> Any:
     )
 
 
-def build_transforms(cfg):
+def build_transforms(cfg: DictConfig):
     if not cfg:
         return transform.Compose([])
     cfg = copy.deepcopy(cfg)
diff --git a/ppsci/data/process/transform/__init__.py b/ppsci/data/process/transform/__init__.py
@@ -12,13 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
 import traceback
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import Tuple
 
 from paddle import vision
 
+if TYPE_CHECKING:
+    from omegaconf import DictConfig
+
 from ppsci.data.process.transform.preprocess import CropData
 from ppsci.data.process.transform.preprocess import FunctionalTransform
 from ppsci.data.process.transform.preprocess import Log1p
@@ -57,7 +63,7 @@ def __call__(self, *data: Tuple[Any, ...]):
         return data
 
 
-def build_transforms(cfg):
+def build_transforms(cfg: DictConfig):
     if not cfg:
         return Compose([])
     cfg = copy.deepcopy(cfg)
diff --git a/ppsci/equation/__init__.py b/ppsci/equation/__init__.py
@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
+from typing import TYPE_CHECKING
 
 from ppsci.equation.fpde import FractionalPoisson
 from ppsci.equation.ide import Volterra
@@ -32,6 +35,10 @@
 from ppsci.utils import logger
 from ppsci.utils import misc
 
+if TYPE_CHECKING:
+    from omegaconf import DictConfig
+
+
 __all__ = [
     "PDE",
     "DETACH_FUNC_NAME",
@@ -52,11 +59,11 @@
 ]
 
 
-def build_equation(cfg):
+def build_equation(cfg: DictConfig):
     """Build equation(s)
 
     Args:
-        cfg (List[DictConfig]): Equation(s) config list.
+        cfg (DictConfig): Equation(s) config list.
 
     Returns:
         Dict[str, Equation]: Equation(s) in dict.
diff --git a/ppsci/geometry/__init__.py b/ppsci/geometry/__init__.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
 
 from ppsci.geometry.geometry import Geometry
@@ -52,11 +54,17 @@
 ]
 
 
-def build_geometry(cfg):
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from omegaconf import DictConfig
+
+
+def build_geometry(cfg: DictConfig):
     """Build geometry(ies)
 
     Args:
-        cfg (List[DictConfig]): Geometry config list.
+        cfg (DictConfig): Geometry config list.
 
     Returns:
         Dict[str, Geometry]: Geometry(ies) in dict.
diff --git a/ppsci/loss/__init__.py b/ppsci/loss/__init__.py
@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
+from typing import TYPE_CHECKING
 
 from ppsci.loss import mtl
 from ppsci.loss.base import Loss
@@ -31,6 +34,10 @@
 from ppsci.loss.mse import MSELossWithL2Decay
 from ppsci.loss.mse import PeriodicMSELoss
 
+if TYPE_CHECKING:
+    from omegaconf import DictConfig
+
+
 __all__ = [
     "Loss",
     "FunctionalLoss",
@@ -51,7 +58,7 @@
 ]
 
 
-def build_loss(cfg):
+def build_loss(cfg: DictConfig):
     """Build loss.
 
     Args:
diff --git a/ppsci/loss/mtl/__init__.py b/ppsci/loss/mtl/__init__.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
 
 from ppsci.loss.mtl.agda import AGDA
@@ -33,7 +35,13 @@
 ]
 
 
-def build_mtl_aggregator(cfg):
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from omegaconf import DictConfig
+
+
+def build_mtl_aggregator(cfg: DictConfig):
     """Build loss aggregator with multi-task learning method.
 
     Args:
diff --git a/ppsci/metric/__init__.py b/ppsci/metric/__init__.py
diff --git a/ppsci/metric/anomaly_coef.py b/ppsci/metric/anomaly_coef.py
diff --git a/ppsci/metric/rmse.py b/ppsci/metric/rmse.py
diff --git a/ppsci/utils/logger.py b/ppsci/utils/logger.py
diff --git a/ppsci/validate/__init__.py b/ppsci/validate/__init__.py
diff --git a/ppsci/visualize/__init__.py b/ppsci/visualize/__init__.py
diff --git a/test/data/test_register_dataset.py b/test/data/test_register_dataset.py
