File size: 8,058 Bytes
338c3e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | import argparse
import json
import importlib.util
import os
import sys
import time
from pathlib import Path
import numpy as np
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import Adam, lr_scheduler
from tqdm import tqdm
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from model import build_model, infer_task_type
import onescience
from onescience.distributed.manager import DistributedManager
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 parse_args():
parser = argparse.ArgumentParser(description="Train CFDBench static or autoregressive models.")
parser.add_argument("--model", default=None, help="Override root.model.name, e.g. ffn, deeponet, fno, auto_ffn.")
return parser.parse_args()
def load_config(model_name=None):
cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
if model_name:
cfg.model.name = model_name
elif os.environ.get("CFDBENCH_MODEL_NAME"):
cfg.model.name = os.environ["CFDBENCH_MODEL_NAME"]
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 checkpoint_path(cfg):
name = cfg.training.get("checkpoint_name", "auto")
if name == "auto":
name = f"{cfg.model.name}.pt"
return PROJECT_ROOT / "weight" / name
def output_path(cfg, task_type):
output_dir = Path(cfg.training.output_dir)
if cfg.training.get("group_by_model", False):
output_dir = output_dir / task_type / cfg.model.name
return output_dir
def select_device(requested, dist):
if requested == "auto":
return dist.device
if requested.startswith("cuda") and not torch.cuda.is_available():
raise RuntimeError(f"Requested device {requested!r}, but CUDA is not available.")
return torch.device(requested)
def dump_json(data, path):
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2)
def load_cfdbench_datapipe_class():
runtime_root = Path(onescience.__file__).resolve().parent
datapipe_file = runtime_root / "datapipes" / "cfd" / "cfdbench.py"
spec = importlib.util.spec_from_file_location("_onescience_cfdbench_datapipe", datapipe_file)
if spec is None or spec.loader is None:
raise ImportError(f"Cannot load CFDBench datapipe from {datapipe_file}")
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module.CFDBenchDatapipe
def mean_scores(score_lists):
return {key: float(np.mean(values)) for key, values in score_lists.items() if values}
def evaluate(model, loader, device, dist, desc="Evaluating"):
model_eval = model.module if hasattr(model, "module") else model
model_eval.eval()
score_lists = {name: [] for name in model_eval.loss_fn.get_score_names()}
with torch.no_grad():
iterator = tqdm(loader, desc=desc, disable=(dist.rank != 0))
for batch in iterator:
batch = {key: value.to(device) for key, value in batch.items()}
outputs = model_eval(**batch)
for key, value in outputs["loss"].items():
score_lists[key].append(float(value.detach().cpu()))
return {"mean": mean_scores(score_lists), "all": score_lists}
def train(cfg, model, datapipe, output_dir, device, dist):
train_cfg = cfg.training
train_loader, train_sampler = datapipe.train_dataloader()
val_loader, _ = datapipe.val_dataloader()
if len(train_loader) == 0:
raise RuntimeError("Training loader is empty. Increase fake_data.num_cases_per_subset or reduce batch_size.")
optimizer = Adam(model.parameters(), lr=train_cfg.lr)
scheduler = lr_scheduler.StepLR(optimizer, step_size=train_cfg.lr_step_size, gamma=train_cfg.lr_gamma)
train_losses = []
best_nmse = float("inf")
for epoch in range(train_cfg.num_epochs):
if train_sampler:
train_sampler.set_epoch(epoch)
model.train()
start = time.time()
iterator = tqdm(train_loader, desc=f"Epoch {epoch}", disable=(dist.rank != 0))
for step, batch in enumerate(iterator):
batch = {key: value.to(device) for key, value in batch.items()}
outputs = model(**batch)
loss = outputs["loss"][train_cfg.loss_name]
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
loss_value = float(loss.detach().cpu())
train_losses.append(loss_value)
if dist.rank == 0 and (step + 1) % train_cfg.log_interval == 0:
iterator.set_postfix({"loss": f"{loss_value:.4e}"})
scheduler.step()
if dist.rank == 0 and (epoch + 1) % train_cfg.eval_interval == 0:
ckpt_dir = output_dir / f"ckpt-{epoch}"
ckpt_dir.mkdir(parents=True, exist_ok=True)
scores = evaluate(model, val_loader, device, dist, desc=f"Val {epoch}")
dump_json(scores, ckpt_dir / "dev_scores.json")
dump_json({"epoch": epoch, "train_loss": train_losses, "seconds": time.time() - start}, ckpt_dir / "scores.json")
nmse = scores["mean"].get(train_cfg.loss_name, float("inf"))
model_to_save = model.module if hasattr(model, "module") else model
torch.save(model_to_save.state_dict(), ckpt_dir / "model.pt")
if nmse <= best_nmse:
best_nmse = nmse
weight_path = checkpoint_path(cfg)
torch.save(model_to_save.state_dict(), weight_path)
print(f"Saved best checkpoint to {weight_path}")
if dist.world_size > 1:
torch.distributed.barrier()
if dist.rank == 0:
dump_json(train_losses, output_dir / "train_losses.json")
def main(required_task_type="static", entry_name="scripts/train.py"):
args = parse_args()
DistributedManager.initialize()
dist = DistributedManager()
cfg = load_config(args.model)
task_type = infer_task_type(cfg.model.name)
if required_task_type is not None and task_type != required_task_type:
expected = "ffn/deeponet" if required_task_type == "static" else "auto_* / resnet / unet / fno"
raise ValueError(
f"{entry_name} is the {required_task_type} entry, but model.name={cfg.model.name!r} is {task_type}. "
f"Use a {expected} model, or run scripts/train_auto.py for autoregressive models."
)
cfg.datapipe.data.task_type = task_type
device = select_device(cfg.training.get("device", "auto"), dist)
output_dir = output_path(cfg, task_type)
if dist.rank == 0:
output_dir.mkdir(parents=True, exist_ok=True)
print(f"Config: {PROJECT_ROOT / 'conf' / 'config.yaml'}")
print(f"Model: {cfg.model.name} ({task_type})")
print(f"Data: {cfg.datapipe.source.data_dir}")
print(f"Output: {output_dir}")
print(f"Checkpoint: {checkpoint_path(cfg)}")
torch.set_num_threads(1)
CFDBenchDatapipe = load_cfdbench_datapipe_class()
datapipe = CFDBenchDatapipe(cfg.datapipe, distributed=(dist.world_size > 1))
model = build_model(cfg).to(device)
if dist.world_size > 1 and "train" in cfg.training.mode:
device_ids = [dist.local_rank] if device.type == "cuda" else None
model = DDP(model, device_ids=device_ids)
if "train" in cfg.training.mode:
train(cfg, model, datapipe, output_dir, device, dist)
if "test" in cfg.training.mode and dist.rank == 0:
scores = evaluate(model, datapipe.test_dataloader(), device, dist, desc="Test")
dump_json(scores, output_dir / "test_scores.json")
print(f"Test scores: {scores['mean']}")
DistributedManager.cleanup()
if __name__ == "__main__":
main()
|