Upload folder using huggingface_hub

README.md

@@ -0,0 +1,145 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- fluid dynamics
+- steady-state laminar flow prediction
+- CFD surrogate modeling
+frameworks: PyTorch
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">DeepCFD</span>
+  </strong>
+</p>
+
+# Model Overview
+DeepCFD is a U-Net-based surrogate model developed by the German Aerospace Center for two-dimensional steady-state laminar flows. It uses geometric information to rapidly predict velocity and pressure fields, accelerating the evaluation of channel flows, flows around obstacles, and aerodynamic designs.
+
+Paper: [DeepCFD: Efficient Steady-State Laminar Flow
+Approximation with Deep Convolutional Neural Networks](https://arxiv.org/abs/2004.08826).
+
+# Model Description
+DeepCFD uses a U-Net deep convolutional architecture that takes a signed distance field (SDF) and flow-region mask as inputs to rapidly predict velocity and pressure fields for two-dimensional, nonuniform, steady-state laminar flows.
+
+## Use Cases
+
+| Use Case | Description |
+| ---------- | ---------------------------------- |
+| Steady-state laminar flow prediction | Approximate velocity and pressure fields in two-dimensional, nonuniform, steady-state laminar flows |
+| Channel-flow simulation | Predict flow distributions around randomly shaped obstacles in a channel |
+| CFD surrogate modeling | Replace conventional CFD workflows such as OpenFOAM to improve inference efficiency |
+| Aerodynamic and fluid-shape optimization | Evaluate large numbers of candidate geometries for low-speed laminar-flow design |
+
+
+# Usage
+
+## 1. OneCode
+
+Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:
+
+[Launch OneCode for one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Setup
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for import checks and small-scale pipeline validation, but full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version for the target cluster, is recommended.
+
+
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/DeepCFD --local_dir ./DeepCFD
+cd DeepCFD
+```
+
+### Set Up the Runtime Environment
+
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Installation with uv is also supported
+pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+```bash
+# Activate Conda first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# Installation with uv is also supported
+pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+The OneScience community provides the `deepcfd` dataset for training. Download it with the command below and verify that the data path in `conf/config.yaml` is configured correctly:
+
+```bash
+modelscope download --dataset OneScience/deepcfd --local_dir ./data
+```
+
+The project dataset can also be [downloaded from Zenodo](https://zenodo.org/record/3666056/files/DeepCFD.zip?download=1).
+
+The dataset contains two files:
+- `dataX.pkl`: Geometric input data for 981 channel-flow samples
+- `dataY.pkl`: Ground-truth CFD solutions for the corresponding samples, computed with the `simpleFOAM` solver
+
+### Training
+
+Single GPU:
+
+- Training parameters are read from the `root.datapipe`, `root.model`, and `root.training` sections of `config/config.yaml`. The default configuration is intended for a minimal smoke test.
+- For full training, set `root.datapipe.source.data_dir` to the actual DeepCFD dataset directory and increase the model size, batch size, and number of epochs as needed.
+
+```bash
+python scripts/train.py
+```
+Multiple GPUs:
+
+```bash
+torchrun --standalone --nproc_per_node=<num_GPUs> scripts/train.py
+```
+By default, training saves the best checkpoint to:
+
+```text
+./weight/best_model.pt
+```
+
+### Model Weights
+
+This repository will provide pretrained DeepCFD weights in the `weights/` directory. The weights will be uploaded soon.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+
+# Official OneScience Resources
+
+| Platform | OneScience Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citations and License
+
+- Original DeepCFD paper: [DeepCFD: Efficient Steady-State Laminar Flow Approximation with Deep Convolutional Neural Networks](https://arxiv.org/abs/2004.08826)
+- This repository retains source attribution and has been adapted for automated execution through OneScience and ModelScope.

config/config.yaml

@@ -0,0 +1,45 @@
+root:
+  datapipe:
+    verbose: false
+    source:
+      data_dir: "./data/deepcfd"
+      data_x_name: "dataX.pkl"
+      data_y_name: "dataY.pkl"
+
+    data:
+      split_ratio: 0.75
+      seed: 0
+
+    dataloader:
+      batch_size: 1
+      num_workers: 0
+
+  model:
+    name: "UNetEx"
+    in_channels: 3
+    out_channels: 3
+    base_channels: 2
+    num_stages: 1
+    kernel_size: 3
+    bilinear: true
+    normtype: "none"
+
+  training:
+    output_dir: "./weight"
+    checkpoint_name: "best_model.pt"
+    num_epochs: 1
+    lr: 0.001
+    weight_decay: 0.005
+    patience: 20
+    eval_interval: 1
+
+  inference:
+    checkpoint_path: "./weight/best_model.pt"
+    result_dir: "./result/deepcfd"
+    num_visualize: 2
+
+  fake_data:
+    num_samples: 4
+    height: 16
+    width: 16
+    seed: 42

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}

model/__init__.py

@@ -0,0 +1,23 @@
+from .unet import UNet
+from .unetex import UNetEx
+
+
+def build_model(config):
+    name = config["name"] if isinstance(config, dict) else config.name
+    model_map = {
+        "UNet": UNet,
+        "UNetEx": UNetEx,
+    }
+    if name not in model_map:
+        raise ValueError(f"Unknown DeepCFD model: {name}")
+
+    get = config.get if isinstance(config, dict) else lambda key, default=None: getattr(config, key, default)
+    return model_map[name](
+        in_channels=get("in_channels"),
+        out_channels=get("out_channels"),
+        base_channels=get("base_channels", 16),
+        num_stages=get("num_stages", 2),
+        bilinear=get("bilinear", True),
+        normtype=get("normtype", "bn"),
+        kernel_size=get("kernel_size", 3),
+    )

model/unet.py

@@ -0,0 +1,44 @@
+import torch.nn as nn
+
+from onescience.modules.decoder.unet_decoder import UNetDecoder2D
+from onescience.modules.encoder.unet_encoder import UNetEncoder2D
+from onescience.modules.head.unet_head import UNetHead2D
+
+
+class UNet(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        base_channels: int = 16,
+        num_stages: int = 2,
+        bilinear: bool = True,
+        normtype: str = "bn",
+        kernel_size: int = 3,
+    ):
+        super().__init__()
+        self.encoder = UNetEncoder2D(
+            in_channels=in_channels,
+            base_channels=base_channels,
+            num_stages=num_stages,
+            bilinear=bilinear,
+            normtype=normtype,
+            kernel_size=kernel_size,
+        )
+        self.decoder = UNetDecoder2D(
+            base_channels=base_channels,
+            num_stages=num_stages,
+            bilinear=bilinear,
+            normtype=normtype,
+            kernel_size=kernel_size,
+        )
+        self.head = UNetHead2D(
+            in_channels=base_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+        )
+
+    def forward(self, x):
+        features = self.encoder(x)
+        decoded = self.decoder(features)
+        return self.head(decoded)

model/unetex.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+
+from onescience.modules.decoder.unet_decoder import UNetDecoder2D
+from onescience.modules.encoder.unet_encoder import UNetEncoder2D
+from onescience.modules.head.unet_head import UNetHead2D
+
+
+class DecoderPath(nn.Module):
+    def __init__(self, base_channels, num_stages, bilinear, normtype, kernel_size):
+        super().__init__()
+        self.decoder = UNetDecoder2D(
+            base_channels=base_channels,
+            num_stages=num_stages,
+            bilinear=bilinear,
+            normtype=normtype,
+            kernel_size=kernel_size,
+        )
+        self.head = UNetHead2D(
+            in_channels=base_channels,
+            out_channels=1,
+            kernel_size=1,
+        )
+
+    def forward(self, features):
+        return self.head(self.decoder(features))
+
+
+class UNetEx(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        base_channels: int = 16,
+        num_stages: int = 2,
+        bilinear: bool = True,
+        normtype: str = "bn",
+        kernel_size: int = 3,
+    ):
+        super().__init__()
+        self.encoder = UNetEncoder2D(
+            in_channels=in_channels,
+            base_channels=base_channels,
+            num_stages=num_stages,
+            bilinear=bilinear,
+            normtype=normtype,
+            kernel_size=kernel_size,
+        )
+        self.decoders = nn.ModuleList(
+            [
+                DecoderPath(base_channels, num_stages, bilinear, normtype, kernel_size)
+                for _ in range(out_channels)
+            ]
+        )
+
+    def forward(self, x):
+        features = self.encoder(x)
+        outputs = [decoder_path(features) for decoder_path in self.decoders]
+        return torch.cat(outputs, dim=1)

scripts/fake_data.py

@@ -0,0 +1,75 @@
+import pickle
+import sys
+from pathlib import Path
+
+import numpy as np
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from onescience.utils.YParams import YParams
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def make_fields(num_samples, height, width, seed):
+    rng = np.random.default_rng(seed)
+    y_axis = np.linspace(-1.0, 1.0, height, dtype=np.float32)
+    x_axis = np.linspace(-1.0, 1.0, width, dtype=np.float32)
+    yy, xx = np.meshgrid(y_axis, x_axis, indexing="ij")
+
+    x = np.empty((num_samples, 3, height, width), dtype=np.float32)
+    target = np.empty_like(x)
+
+    for i in range(num_samples):
+        cx = 0.2 * np.sin(i)
+        cy = 0.2 * np.cos(i * 0.7)
+        radius = 0.25 + 0.03 * (i % 3)
+        distance = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2) - radius
+        obstacle = (distance < 0.0).astype(np.float32)
+        inlet = np.ones_like(xx, dtype=np.float32) * (1.0 + 0.05 * i)
+
+        x[i, 0] = obstacle
+        x[i, 1] = distance
+        x[i, 2] = inlet
+
+        ux = inlet * (1.0 - obstacle) + 0.1 * np.sin(np.pi * yy)
+        uy = -0.2 * yy * np.exp(-4.0 * np.maximum(distance, 0.0))
+        pressure = 0.5 * (1.0 - xx) + 0.1 * obstacle
+
+        noise = 0.005 * rng.standard_normal((3, height, width)).astype(np.float32)
+        target[i, 0] = ux
+        target[i, 1] = uy
+        target[i, 2] = pressure
+        target[i] += noise
+
+    return x, target
+
+
+def main():
+    cfg = YParams(str(PROJECT_ROOT / "config" / "config.yaml"), "root")
+    data_dir = resolve_path(cfg.datapipe.source.data_dir)
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    x, y = make_fields(
+        num_samples=cfg.fake_data.num_samples,
+        height=cfg.fake_data.height,
+        width=cfg.fake_data.width,
+        seed=cfg.fake_data.seed,
+    )
+
+    with open(data_dir / cfg.datapipe.source.data_x_name, "wb") as f:
+        pickle.dump(x, f)
+    with open(data_dir / cfg.datapipe.source.data_y_name, "wb") as f:
+        pickle.dump(y, f)
+
+    print(f"Fake DeepCFD data written to {data_dir}")
+    print(f"x shape={x.shape}, y shape={y.shape}, dtype={x.dtype}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,103 @@
+import os
+import sys
+import importlib.util
+from pathlib import Path
+
+os.environ.setdefault("MPLBACKEND", "Agg")
+
+import numpy as np
+import torch
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model import build_model
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+from onescience.utils.deepcfd.functions import visualize
+import onescience
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def load_config():
+    cfg = YParams(str(PROJECT_ROOT / "config" / "config.yaml"), "root")
+    cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
+    cfg.inference.checkpoint_path = str(resolve_path(cfg.inference.checkpoint_path))
+    cfg.inference.result_dir = str(resolve_path(cfg.inference.result_dir))
+    return cfg
+
+
+def load_deepcfd_datapipe_class():
+    runtime_root = Path(onescience.__file__).resolve().parent
+    datapipe_file = runtime_root / "datapipes" / "cfd" / "deepcfd.py"
+    spec = importlib.util.spec_from_file_location("_onescience_deepcfd_datapipe", datapipe_file)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Cannot load DeepCFD datapipe from {datapipe_file}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module.DeepCFDDatapipe
+
+
+def main():
+    DistributedManager.initialize()
+    dist = DistributedManager()
+    device = dist.device
+    cfg = load_config()
+    DeepCFDDatapipe = load_deepcfd_datapipe_class()
+
+    checkpoint_path = Path(cfg.inference.checkpoint_path)
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model_config = checkpoint.get("config", cfg.model.to_dict())
+    model = build_model(model_config).to(device)
+    model.load_state_dict(checkpoint["model_state"])
+    model.eval()
+
+    datapipe = DeepCFDDatapipe(cfg.datapipe, distributed=False)
+    test_loader, _ = datapipe.test_dataloader()
+    batch = next(iter(test_loader))
+
+    x = batch["x"].to(device)
+    y = batch["y"].to(device)
+    with torch.no_grad():
+        out = model(x)
+
+    error = torch.abs(out.cpu() - y.cpu())
+    mse = torch.mean((out.cpu() - y.cpu()) ** 2, dim=(0, 2, 3)).numpy()
+    mae = torch.mean(error, dim=(0, 2, 3)).numpy()
+
+    result_dir = Path(cfg.inference.result_dir)
+    vis_dir = result_dir / "vis_results"
+    pred_dir = result_dir / "predictions"
+    vis_dir.mkdir(parents=True, exist_ok=True)
+    pred_dir.mkdir(parents=True, exist_ok=True)
+
+    np.save(pred_dir / "prediction_batch.npy", out.cpu().numpy())
+    np.save(pred_dir / "target_batch.npy", y.cpu().numpy())
+    np.save(pred_dir / "absolute_error_batch.npy", error.numpy())
+
+    y_np = y.cpu().numpy()
+    out_np = out.cpu().numpy()
+    err_np = error.numpy()
+    for i in range(min(cfg.inference.num_visualize, x.shape[0])):
+        visualize(y_np, out_np, err_np, i, save_dir=str(vis_dir))
+
+    if dist.rank == 0:
+        print(f"Checkpoint: {checkpoint_path}")
+        print(f"MSE per channel [Ux, Uy, p]: {mse}")
+        print(f"MAE per channel [Ux, Uy, p]: {mae}")
+        print(f"Results saved to {result_dir}")
+
+    dist.cleanup()
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,41 @@
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from onescience.utils.YParams import YParams
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def main():
+    cfg = YParams(str(PROJECT_ROOT / "config" / "config.yaml"), "root")
+    checkpoint_path = resolve_path(cfg.inference.checkpoint_path)
+    pred_dir = resolve_path(cfg.inference.result_dir) / "predictions"
+
+    if checkpoint_path.exists():
+        ckpt = torch.load(checkpoint_path, map_location="cpu")
+        print(f"Checkpoint: {checkpoint_path}")
+        print(f"Epoch: {ckpt.get('epoch')}, val_loss: {ckpt.get('val_loss')}")
+        print(f"Model config: {ckpt.get('config')}")
+    else:
+        print(f"Checkpoint not found: {checkpoint_path}")
+
+    pred_path = pred_dir / "prediction_batch.npy"
+    if pred_path.exists():
+        pred = np.load(pred_path)
+        print(f"Prediction batch: shape={pred.shape}, dtype={pred.dtype}, mean={pred.mean():.6f}")
+    else:
+        print(f"Prediction batch not found: {pred_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,170 @@
+import sys
+import importlib.util
+from pathlib import Path
+
+import torch
+from torch.nn.parallel import DistributedDataParallel as DDP
+from tqdm import tqdm
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model import build_model
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+import onescience
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def load_config():
+    cfg = YParams(str(PROJECT_ROOT / "config" / "config.yaml"), "root")
+    cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
+    cfg.training.output_dir = str(resolve_path(cfg.training.output_dir))
+    return cfg
+
+
+def load_deepcfd_datapipe_class():
+    runtime_root = Path(onescience.__file__).resolve().parent
+    datapipe_file = runtime_root / "datapipes" / "cfd" / "deepcfd.py"
+    spec = importlib.util.spec_from_file_location("_onescience_deepcfd_datapipe", datapipe_file)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Cannot load DeepCFD datapipe from {datapipe_file}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module.DeepCFDDatapipe
+
+
+def loss_func(output, target, weights):
+    lossu = (output[:, 0] - target[:, 0]) ** 2
+    lossv = (output[:, 1] - target[:, 1]) ** 2
+    lossp = torch.abs(output[:, 2] - target[:, 2])
+    loss_stack = torch.stack([lossu, lossv, lossp], dim=1)
+    return torch.sum(loss_stack / weights)
+
+
+def evaluate(model, loader, device, weights, dist):
+    model.eval()
+    total_loss = 0.0
+    total_ux_mse = 0.0
+    total_uy_mse = 0.0
+    total_p_mse = 0.0
+    num_batches = 0
+
+    with torch.no_grad():
+        iterator = tqdm(loader, desc="Evaluating", disable=(dist.rank != 0))
+        for batch in iterator:
+            x = batch["x"].to(device)
+            y = batch["y"].to(device)
+            output = model(x)
+
+            total_loss += loss_func(output, y, weights).item()
+            total_ux_mse += torch.sum((output[:, 0] - y[:, 0]) ** 2).item()
+            total_uy_mse += torch.sum((output[:, 1] - y[:, 1]) ** 2).item()
+            total_p_mse += torch.sum((output[:, 2] - y[:, 2]) ** 2).item()
+            num_batches += 1
+
+    if num_batches == 0:
+        raise RuntimeError("Evaluation loader is empty. Check split_ratio and dataset size.")
+    return total_loss / num_batches, total_ux_mse, total_uy_mse, total_p_mse
+
+
+def main():
+    DistributedManager.initialize()
+    dist = DistributedManager()
+    device = dist.device
+    cfg = load_config()
+    DeepCFDDatapipe = load_deepcfd_datapipe_class()
+
+    output_dir = Path(cfg.training.output_dir)
+    if dist.rank == 0:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        print(f"Config: {PROJECT_ROOT / 'config' / 'config.yaml'}")
+        print(f"Data: {cfg.datapipe.source.data_dir}")
+        print(f"Checkpoint directory: {output_dir}")
+
+    datapipe = DeepCFDDatapipe(cfg.datapipe, distributed=(dist.world_size > 1))
+    train_loader, train_sampler = datapipe.train_dataloader()
+    test_loader, _ = datapipe.test_dataloader()
+    loss_weights = datapipe.get_loss_weights().to(device)
+
+    model = build_model(cfg.model).to(device)
+    if dist.world_size > 1:
+        device_ids = [dist.local_rank] if device.type == "cuda" else None
+        model = DDP(model, device_ids=device_ids)
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=cfg.training.lr,
+        weight_decay=cfg.training.weight_decay,
+    )
+
+    best_val_loss = float("inf")
+    patience_counter = 0
+
+    for epoch in range(cfg.training.num_epochs):
+        if train_sampler:
+            train_sampler.set_epoch(epoch)
+
+        model.train()
+        train_loss = 0.0
+        iterator = tqdm(train_loader, desc=f"Epoch {epoch}", disable=(dist.rank != 0))
+
+        for batch in iterator:
+            x = batch["x"].to(device)
+            y = batch["y"].to(device)
+
+            optimizer.zero_grad(set_to_none=True)
+            output = model(x)
+            loss = loss_func(output, y, loss_weights)
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss.item()
+            if dist.rank == 0:
+                iterator.set_postfix({"loss": f"{loss.item():.4e}"})
+
+        if len(train_loader) == 0:
+            raise RuntimeError("Training loader is empty. Check split_ratio and dataset size.")
+        avg_train_loss = train_loss / len(train_loader)
+
+        if (epoch + 1) % cfg.training.eval_interval == 0:
+            val_loss, ux_err, uy_err, p_err = evaluate(model, test_loader, device, loss_weights, dist)
+
+            if dist.rank == 0:
+                print(f"Epoch {epoch} | Train Loss: {avg_train_loss:.4e} | Val Loss: {val_loss:.4e}")
+                print(f"Metrics (Sum Sq Err): Ux={ux_err:.2e}, Uy={uy_err:.2e}, P={p_err:.2e}")
+
+                if val_loss < best_val_loss:
+                    best_val_loss = val_loss
+                    patience_counter = 0
+                    model_to_save = model.module if hasattr(model, "module") else model
+                    ckpt = {
+                        "model_state": model_to_save.state_dict(),
+                        "config": cfg.model.to_dict(),
+                        "epoch": epoch,
+                        "val_loss": val_loss,
+                    }
+                    torch.save(ckpt, output_dir / cfg.training.checkpoint_name)
+                    print(f"Saved best model to {output_dir / cfg.training.checkpoint_name}")
+                else:
+                    patience_counter += 1
+
+            stop_flag = torch.tensor([0], device=device)
+            if dist.rank == 0 and patience_counter >= cfg.training.patience:
+                stop_flag += 1
+            if dist.world_size > 1:
+                torch.distributed.broadcast(stop_flag, src=0)
+            if stop_flag.item() > 0:
+                break
+
+    dist.cleanup()
+
+
+if __name__ == "__main__":
+    main()

weight/.gitkeep

@@ -0,0 +1 @@
+