huangfu/photonAI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修改脚本结构MOE

Browse Source
This commit is contained in:
huangfu
2026-04-19 18:33:36 +08:00
parent 745868a456
commit 6693f3916c
13 changed files with 1675 additions and 1096 deletions
Download Patch File Download Diff File Expand all files Collapse all files

949
MZM_MoE_PINN_Model.ipynb Normal file
View File

File diff suppressed because one or more lines are too long

203
README.md
View File

@@ -1,163 +1,146 @@
# MZM 器件性能 MLP 回归基线（PyTorch）
# MZM MoE PINN（PyTorch）
本项目实现一个**多输出回归**基线模型：用 **8 个器件/偏置参数**预测 **3 个射频性能指标**。代码面向科研复现：可配置 YAML、固定随机种子、训练集拟合标准化器、完整日志与可视化产物。
本仓库已从原先的 MLP 基线迁移为 **Mixture-of-Experts + Physics-Informed Neural Network** 训练流程，目标是尽量对齐 `MZM_MoE_PINN_Model.ipynb` 的训练行为，同时保留仓库化的命令行入口、结果目录和复现产物。
## 项目简介
## 当前训练管线
- **任务类型**：监督学习，多输出回归（非分类）。
- **输入（8 维）**：工艺与偏置相关参数。
- **输出（3 维）**：`BW_3dB`、`IL`、`V_pi`。
- **模型**：原生 PyTorch MLP，可选 BatchNorm / Dropout / 残差（同维时相加）。
- **损失**：默认在**标准化后的输出空间**使用加权 `SmoothL1Loss`（Huber）；可选加权 MSE。
- **v1 目标**：先把数据清洗、划分、训练、评估、日志与可视化流程跑通；**不引入 physics loss**。
- **输入**：8 个器件/偏置参数
- **输出**：`BW_3dB`、`IL`、`V_pi`
- **模型**：MoE，包含多个专家网络与一个 gating 网络
- **数据清洗**：保留 `V_pi < 500`
- **数据划分**：`train_test_split(test_size=0.1, random_state=123)`
- **标准化**：
- `X` 使用一个 `StandardScaler`
- `Y` 的三个目标分别使用独立的 `StandardScaler`
- **损失**：
- 数据项：标准化空间 `MSE`
- 物理项：`dBW/dL <= 0`、`dIL/dL >= 0`、`d(V_pi*L)/dL ~= 0`、`d2BW/dL2` 平滑项
- **优化器**：`AdamW(lr=1e-3, weight_decay=0.05, betas=(0.9, 0.999))`
- **训练方式**：固定 `100` epoch，无早停；每轮计算全量 train/test MSE
默认物理约束权重来自 `best_hyperparams.json`：
```json
{
"lambda_bw_mon": 0.0,
"lambda_IL_mon": 0.3,
"lambda_vpiL": 0.005,
"lambda_smooth": 0.1
}
```
## 数据格式
数据为 **txt 或 csv**，每行 **11 个逗号分隔的浮点数**，无表头（txt）或表头与下列字段一致（csv）。
数据文件为 11 列逗号分隔浮点数，列含义如下：
| 顺序 | 列名 | 含义 | 作为 |
| --- | --- | --- | --- |
| 1 | `PN_offset` | PN 偏移 | 输入 |
| 2 | `Bias_V` | 偏置电压 | 输入 |
| 3 | `Core_width` | 芯区宽度 | 输入 |
| 4 | `P+_width` | P+ 区宽度 | 输入 |
| 5 | `N+_width` | N+ 区宽度 | 输入 |
| 6 | `P_width` | P 区宽度 | 输入 |
| 7 | `N_width` | N 区宽度 | 输入 |
| 8 | `Phase_length` | 相位区长度 | 输入 |
| 9 | `BW_3dB` | 3 dB 带宽 | 目标 |
| 10 | `IL` | 插入损耗 | 目标 |
| 11 | `V_pi` | 半波电压 | 目标 |
| 顺序 | 列名 | 作为 |
| --- | --- | --- |
| 1 | `PN_offset` | 输入 |
| 2 | `Bias_V` | 输入 |
| 3 | `Core_width` | 输入 |
| 4 | `P+_width` | 输入 |
| 5 | `N+_width` | 输入 |
| 6 | `P_width` | 输入 |
| 7 | `N_width` | 输入 |
| 8 | `Phase_length` | 输入 |
| 9 | `BW_3dB` | 输出 |
| 10 | `IL` | 输出 |
| 11 | `V_pi` | 输出 |
- 自动忽略空行与行首行尾空格。
- 每行必须恰好 **11 列**；也支持仿真导出的 **整行方括号** 写法，例如 `[a, b, ..., k]`（与无括号的 `a, b, ..., k` 等价）。
- 否则整文件解析失败并给出错误行号提示。
支持两种文本格式：
## TXT 数据清洗流程（以 V_pi 为准）
```text
a,b,c,...,k
[a, b, c, ..., k]
```
本仓库约定：**txt 每行从左到右第 11 个逗号分隔浮点数**即半波电压 **`V_pi`**（与表头列名一致）。清洗时以该列为**物理可信区间**的主门控，避免异常仿真/标注污染训练。
建议按以下顺序理解流水线（与 `src/preprocess.py` 中 `clean_dataframe` 实现一致）：
1. **解析与建表**：读取 txt → 校验每行 11 列 → 转为 `float` → 构建 `DataFrame`（最后一列为 `V_pi`）。
2. **（可选）去重**：`remove_duplicate_rows: true` 时删除 11 列完全相同的重复行。
3. **V_pi 区间门控（主清洗）**：默认启用 `filter_v_pi_range: true`，仅保留
`v_pi_min <= V_pi <= v_pi_max`（默认 **`[0, 500]`**）。**区间之外整行剔除**。
该步骤专门针对「以最后一列 `V_pi` 为正常范围」的需求。
4. **（可选）严格正电压**：`remove_nonpositive_vpi: true` 时，在区间过滤之后再删除 `V_pi <= 0`（若需保留 `V_pi = 0` 且仍在 `[0,500]` 内，请保持为 `false`）。
5. **后续步骤**：默认采用**按 8 个输入字段分组**的 train/val/test 切分，避免「同输入异输出」同时落入不同集合；再按 `split_stratify_target`（默认 `V_pi`）做组级近似分层；之后才做（可选）训练集离群策略与仅在训练集上拟合 `StandardScaler`。
清洗前会在日志与 `data_report.md` 中报告：给定 `[v_pi_min, v_pi_max]` 下 **`V_pi` 越界行数**、重复样本、同输入异输出等统计，便于核对。
## 环境要求
- Python **3.10+**（已在 3.13 下通过冒烟测试）。
- 推荐使用虚拟环境。
### 安装依赖
## 安装
```bash
cd /path/to/photonAI
python -m venv .venv
source .venv/bin/activate # Windows 使用 .venv\Scripts\activate
source .venv/bin/activate
pip install -U pip
pip install -r requirements.txt
```
## 放置数据
1. 将原始 txt（例如仓库根目录下的 `Sim_MZM_dataset.txt`）复制或软链接到 `data/dataset.txt`。
2. 或在 `configs/default.yaml` 中修改 `data_path` 为绝对路径或相对项目根目录的路径。
若路径不存在，程序会给出明确报错，不会静默失败。
## 训练
```bash
python -m src.main train --config configs/default.yaml
```
或使用脚本：
训练完成后会在 `results/run_时间戳/` 下生成：
```bash
bash scripts/train.sh
```
- `config_snapshot.yaml`
- `data_report.md`
- `data_stats.csv`
- `cleaning_meta.json`
- `split_indices.json`
- `x_scaler.pkl`
- `y_scalers.pkl`
- `train_log.csv`
- `checkpoints/best.pt`
- `checkpoints/last.pt`
- `metrics.csv`
- `summary.json`
- `summary.md`
- `test_predictions.csv`
- `figures/*.png`
训练会在 `results/run_时间戳/` 下生成：
说明：
- `config_snapshot.yaml`：本次运行配置快照。
- `split_indices.json`：对**清洗后**样本行的 train/val/test 索引，便于 `eval` 完全复现划分。
- `x_scaler.pkl` / `y_scaler.pkl`：`StandardScaler`，推理阶段用于反标准化。
- `data_report.md` / `data_stats.csv`：数据统计与清洗说明。
- `cleaning_meta.json`：清洗与划分元信息。
- `train_log.csv`：逐 epoch 的 train/val loss 与学习率。
- `checkpoints/best.pt`、`checkpoints/last.pt`：最优与最后一轮权重。
- 训练结束后：`metrics.csv`、`summary.json`、`summary.md`、`test_predictions.csv`、`figures/*.png`。
- `train_log.csv` 记录每轮的全量 `train_loss` / `test_loss`
- `summary.*` 与 `metrics.csv` 中的 `loss` 为**标准化空间 MSE**
- 物理空间指标仍输出 `MAE / RMSE / R²`
**说明（损失列）**：`metrics.csv` / `summary.*` 中的 `loss` 与 `*_loss` 均在**标准化输出空间**按训练准则（Huber / 加权 MSE）计算；物理量空间以 **MAE / RMSE / R²** 为主指标。
## 评估
## 评估（复现划分与 scaler）
在**同一数据文件**与 `config_snapshot.yaml` 前提下，可仅运行评估：
按训练时保存的切分索引与 scaler 重算 train/test 指标：
```bash
python -m src.main eval --config configs/default.yaml --run-dir results/run_YYYYMMDD_HHMMSS
```
若不指定 `--run-dir`，将在 `configs/default.yaml` 的 `output_dir`（默认 `results`）下自动选择**最近修改时间**的 `run_*` 目录。
```bash
bash scripts/eval.sh --run-dir results/run_某次训练
```
## 推理
输入文件需包含上述 **8 个输入列**（csv 带表头，或 8 列无表头 txt）。
输入文件需包含 8 个输入列（csv 带表头，或 8 列 txt）：
```bash
python -m src.main infer --config configs/default.yaml --input path/to/inputs.csv --output path/to/preds.csv
```
脚本封装：
输出列为原始 8 个输入 + `pred_BW_3dB`、`pred_IL`、`pred_V_pi`。
```bash
bash scripts/infer.sh path/to/inputs.csv --run-dir results/run_某次训练 --output preds.csv
```
## 默认配置
输出列为 8 个输入 + `pred_BW_3dB`、`pred_IL`、`pred_V_pi`（**物理量空间**，已反标准化）。
`configs/default.yaml` 目前对应 notebook 风格的默认 MoE PINN 参数：
- `data.test_size: 0.1`
- `data.random_state: 123`
- `data.filter_v_pi_max: 500.0`
- `model.hidden_dims: [64, 128, 64]`
- `model.n_experts: 60`
- `model.gating_hidden: 8`
- `model.dropout_rate: 0.0`
- `model.use_bn: true`
- `optimizer.lr: 0.001`
- `optimizer.weight_decay: 0.05`
- `training.batch_size: 128`
- `training.epochs: 100`
- `physics.*` 默认由 `best_hyperparams.json` 提供，再由 YAML 显式值覆盖
## 测试
```bash
pip install pytest
pytest -q tests/test_smoke.py
PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest -q tests/test_smoke.py
```
## 配置说明（`configs/default.yaml`）
## 说明
主要字段：
- **数据与清洗**：`data_path`、`remove_duplicate_rows`、**`filter_v_pi_range` / `v_pi_min` / `v_pi_max`**（默认按 **`V_pi ∈ [0, 500]`** 剔除越界行，对应 txt **第 11 列**）、`remove_nonpositive_vpi`、`outlier_strategy`（`none` / `iqr` / `zscore` / `quantile_clip`）及 `outlier_apply_to`（`targets` / `all`）。
- **划分**：`split_ratios`、`random_seed`、`split_mode`、`split_stratify_target`、`split_stratify_bins`。默认 `grouped_stratified`：先按 8 维输入分组，再按指定目标（默认 `V_pi`）做组级近似分层；也可切回 `random`。**仅在训练子集**上拟合标准化器；离群阈值（若启用）也在训练子集上统计。
- **模型**：`hidden_dims`、`batchnorm`、`dropout`、`residual`。
- **训练**：`AdamW`、`lr`、`weight_decay`、`batch_size`、`epochs`、早停 `early_stopping_patience`。
- **调度器**：`cosine`（默认）或 `plateau`。
- **损失**：`huber`（默认）或 `weighted_mse`，`target_weights` 长度须为 3。
默认策略刻意**不删除**仅因统计极端的样本（`outlier_strategy: none`），但在报告中给出极端值计数；**默认以 `V_pi` 物理区间 `[0,500]` 删除越界行**；`remove_nonpositive_vpi` 默认为 `false`，以便与「0 属于合法下界」一致，需要时可改为 `true`。
## 项目结构
```text
.
├── README.md
├── requirements.txt
├── .gitignore
├── configs
│ └── default.yaml
├── data
├── reports
├── results
- 现在的主流程优先保证与 notebook 的 **数据切分、标准化、模型结构、物理损失和训练循环** 一致。
- 为了适配仓库化使用，仍保留了 `train / eval / infer` CLI 与 `run_*` 结果目录结构。
- 旧的 MLP baseline 文档与配置已不再是当前默认路径。
├── scripts
│ ├── train.sh
│ ├── eval.sh

11
best_hyperparams.json Normal file
View File

@@ -0,0 +1,11 @@
{
"best_config": {
"lambda_bw_mon": 0.0,
"lambda_IL_mon": 0.3,
"lambda_vpiL": 0.005,
"lambda_smooth": 0.1
},
"final_train_loss": 0.011141298338770866,
"final_test_loss": 0.013557782396674156,
"n_params": 619503
}

73
configs/default.yaml
View File

@@ -1,70 +1,39 @@
# 默认配置：MZM MLP 多输出回归基线
# 将数据 txt 放到 data/ 下并修改 data_path，或保持路径指向你的文件
# 默认配置：MZM MoE PINN（按 notebook 迁移）
data_path: data/dataset.txt
best_hyperparams_path: best_hyperparams.json
split_ratios: [0.7, 0.15, 0.15] # train, val, test；可改为 [0.8, 0.1, 0.1]
random_seed: 42
# 切分策略：按 8 维输入分组，避免“同输入异输出”跨集合泄漏；再按目标分桶近似分层
split_mode: grouped_stratified # grouped_stratified | random
split_stratify_target: V_pi
split_stratify_bins: 10
remove_duplicate_rows: true
# 异常值处理策略：none | iqr | zscore | quantile_clip
# 默认仅报告极端值，不删除；物理上不可信的 V_pi 由下方区间门控剔除
outlier_strategy: none
# 启用非 none 策略时，在训练子集上拟合阈值；iqr/zscore 仅删训练集离群行；quantile_clip 按训练分位数 winsorize
outlier_apply_to: targets # targets | all
outlier_config:
iqr_k: 1.5
zscore_threshold: 4.0
quantile_lower: 0.001
quantile_upper: 0.999
# 以 txt 第 11 列（列名 V_pi）为物理门控：仅保留闭区间 [v_pi_min, v_pi_max] 内样本
filter_v_pi_range: true
v_pi_min: 0.0
v_pi_max: 500.0
# 在区间过滤之后，是否再剔除 V_pi<=0；若需保留 V_pi=0（仍在 [0,500] 内），请设为 false
remove_nonpositive_vpi: false
data:
test_size: 0.1
random_state: 123
filter_v_pi_max: 500.0
model:
input_dim: 8
hidden_dims: [200, 300, 350, 300, 200]
output_dim: 3
batchnorm: false
# 温和 dropout，实测略优于全 0（见 results/run_20260419_163305）
dropout: 0.05
residual: false
hidden_dims: [64, 128, 64]
n_experts: 60
gating_hidden: 8
dropout_rate: 0.0
use_bn: true
activation: relu
optimizer:
name: adamw
lr: 0.001
weight_decay: 0.0001
scheduler:
type: cosine # cosine | plateau
plateau_factor: 0.5
plateau_patience: 10
plateau_min_lr: 1.0e-6
weight_decay: 0.05
betas: [0.9, 0.999]
training:
batch_size: 128
epochs: 300
early_stopping_patience: 30
epochs: 100
num_workers: 0
loss:
type: huber # huber | weighted_mse
huber_delta: 1.0
# BW_3dB, IL, V_pi；略加重 V_pi 以小幅提升其测试 R²
target_weights: [1.0, 1.0, 1.2]
# 默认会先从 best_hyperparams.json 读取这些系数，再用此处显式值覆盖
physics:
lambda_bw_mon: 0.0
lambda_IL_mon: 0.3
lambda_vpiL: 0.005
lambda_smooth: 0.1
# 总输出目录；每次训练会在其下创建 run_时间戳/
output_dir: results
# 评估/推理时若未指定 run_dir，可填最近一次 run 的路径（可选）
last_run_dir: null

203
src/config.py
View File

@@ -2,6 +2,7 @@
from __future__ import annotations
import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, List, Optional
@@ -9,137 +10,118 @@ from typing import Any, List, Optional
import yaml
@dataclass
class DataConfig:
test_size: float = 0.1
random_state: int = 123
filter_v_pi_max: float = 500.0
@dataclass
class ModelConfig:
input_dim: int = 8
hidden_dims: List[int] = field(default_factory=lambda: [200, 300, 350, 300, 200])
output_dim: int = 3
batchnorm: bool = False
dropout: float = 0.0
residual: bool = False
hidden_dims: List[int] = field(default_factory=lambda: [64, 128, 64])
n_experts: int = 60
gating_hidden: int = 8
dropout_rate: float = 0.0
use_bn: bool = True
activation: str = "relu"
@dataclass
class OptimizerConfig:
name: str = "adamw"
lr: float = 1e-3
weight_decay: float = 1e-4
@dataclass
class SchedulerConfig:
type: str = "cosine" # cosine | plateau
plateau_factor: float = 0.5
plateau_patience: int = 10
plateau_min_lr: float = 1e-6
weight_decay: float = 0.05
betas: List[float] = field(default_factory=lambda: [0.9, 0.999])
@dataclass
class TrainingConfig:
batch_size: int = 128
epochs: int = 300
early_stopping_patience: int = 30
epochs: int = 100
num_workers: int = 0
@dataclass
class LossConfig:
type: str = "huber" # huber | weighted_mse
huber_delta: float = 1.0
target_weights: List[float] = field(default_factory=lambda: [1.0, 1.0, 1.0])
@dataclass
class OutlierConfig:
iqr_k: float = 1.5
zscore_threshold: float = 4.0
quantile_lower: float = 0.001
quantile_upper: float = 0.999
class PhysicsConfig:
lambda_bw_mon: float = 0.0
lambda_IL_mon: float = 0.3
lambda_vpiL: float = 0.005
lambda_smooth: float = 0.1
@dataclass
class AppConfig:
data_path: str
split_ratios: List[float]
random_seed: int
split_mode: str
split_stratify_target: str
split_stratify_bins: int
remove_duplicate_rows: bool
outlier_strategy: str
outlier_config: OutlierConfig
outlier_apply_to: str # targets | all
remove_nonpositive_vpi: bool
filter_v_pi_range: bool
v_pi_min: float
v_pi_max: float
data: DataConfig
model: ModelConfig
optimizer: OptimizerConfig
scheduler: SchedulerConfig
training: TrainingConfig
loss: LossConfig
physics: PhysicsConfig
output_dir: str
best_hyperparams_path: Optional[str] = None
last_run_dir: Optional[str] = None
@staticmethod
def from_dict(raw: dict[str, Any]) -> "AppConfig":
m = raw.get("model", {})
o = raw.get("optimizer", {})
s = raw.get("scheduler", {})
t = raw.get("training", {})
l = raw.get("loss", {})
oc = raw.get("outlier_config", {})
def from_dict(raw: dict[str, Any], cfg_dir: Path) -> "AppConfig":
data_raw = raw.get("data", {})
model_raw = raw.get("model", {})
optimizer_raw = raw.get("optimizer", {})
training_raw = raw.get("training", {})
physics_raw = raw.get("physics", {})
best_hyperparams_path = raw.get("best_hyperparams_path")
if best_hyperparams_path:
hp_path = Path(best_hyperparams_path)
if not hp_path.is_absolute():
cand = (cfg_dir / hp_path).resolve()
if cand.is_file():
hp_path = cand
else:
hp_path = (cfg_dir.parent / hp_path).resolve()
with hp_path.open("r", encoding="utf-8") as f:
hp_raw = json.load(f)
physics_raw = {**hp_raw.get("best_config", {}), **physics_raw}
best_hyperparams_path = str(hp_path)
return AppConfig(
data_path=str(raw["data_path"]),
split_ratios=list(raw["split_ratios"]),
random_seed=int(raw["random_seed"]),
split_mode=str(raw.get("split_mode", "grouped_stratified")),
split_stratify_target=str(raw.get("split_stratify_target", "V_pi")),
split_stratify_bins=int(raw.get("split_stratify_bins", 10)),
remove_duplicate_rows=bool(raw["remove_duplicate_rows"]),
outlier_strategy=str(raw.get("outlier_strategy", "none")),
outlier_config=OutlierConfig(
iqr_k=float(oc.get("iqr_k", 1.5)),
zscore_threshold=float(oc.get("zscore_threshold", 4.0)),
quantile_lower=float(oc.get("quantile_lower", 0.001)),
quantile_upper=float(oc.get("quantile_upper", 0.999)),
data=DataConfig(
test_size=float(data_raw.get("test_size", raw.get("test_size", 0.1))),
random_state=int(data_raw.get("random_state", raw.get("random_seed", 123))),
filter_v_pi_max=float(
data_raw.get("filter_v_pi_max", raw.get("v_pi_max", 500.0))
),
),
outlier_apply_to=str(raw.get("outlier_apply_to", "targets")),
remove_nonpositive_vpi=bool(raw.get("remove_nonpositive_vpi", False)),
filter_v_pi_range=bool(raw.get("filter_v_pi_range", True)),
v_pi_min=float(raw.get("v_pi_min", 0.0)),
v_pi_max=float(raw.get("v_pi_max", 500.0)),
model=ModelConfig(
input_dim=int(m.get("input_dim", 8)),
hidden_dims=list(m.get("hidden_dims", [200, 300, 350, 300, 200])),
output_dim=int(m.get("output_dim", 3)),
batchnorm=bool(m.get("batchnorm", False)),
dropout=float(m.get("dropout", 0.0)),
residual=bool(m.get("residual", False)),
input_dim=int(model_raw.get("input_dim", 8)),
output_dim=int(model_raw.get("output_dim", 3)),
hidden_dims=list(model_raw.get("hidden_dims", [64, 128, 64])),
n_experts=int(model_raw.get("n_experts", 60)),
gating_hidden=int(model_raw.get("gating_hidden", 8)),
dropout_rate=float(model_raw.get("dropout_rate", 0.0)),
use_bn=bool(model_raw.get("use_bn", True)),
activation=str(model_raw.get("activation", "relu")),
),
optimizer=OptimizerConfig(
name=str(o.get("name", "adamw")),
lr=float(o.get("lr", 1e-3)),
weight_decay=float(o.get("weight_decay", 1e-4)),
),
scheduler=SchedulerConfig(
type=str(s.get("type", "cosine")),
plateau_factor=float(s.get("plateau_factor", 0.5)),
plateau_patience=int(s.get("plateau_patience", 10)),
plateau_min_lr=float(s.get("plateau_min_lr", 1e-6)),
lr=float(optimizer_raw.get("lr", 1e-3)),
weight_decay=float(optimizer_raw.get("weight_decay", 0.05)),
betas=[float(x) for x in optimizer_raw.get("betas", [0.9, 0.999])],
),
training=TrainingConfig(
batch_size=int(t.get("batch_size", 128)),
epochs=int(t.get("epochs", 300)),
early_stopping_patience=int(t.get("early_stopping_patience", 30)),
num_workers=int(t.get("num_workers", 0)),
batch_size=int(training_raw.get("batch_size", 128)),
epochs=int(training_raw.get("epochs", 100)),
num_workers=int(training_raw.get("num_workers", 0)),
),
loss=LossConfig(
type=str(l.get("type", "huber")),
huber_delta=float(l.get("huber_delta", 1.0)),
target_weights=[float(x) for x in l.get("target_weights", [1.0, 1.0, 1.0])],
physics=PhysicsConfig(
lambda_bw_mon=float(physics_raw.get("lambda_bw_mon", 0.0)),
lambda_IL_mon=float(physics_raw.get("lambda_IL_mon", 0.3)),
lambda_vpiL=float(physics_raw.get("lambda_vpiL", 0.005)),
lambda_smooth=float(physics_raw.get("lambda_smooth", 0.1)),
),
output_dir=str(raw.get("output_dir", "results")),
best_hyperparams_path=best_hyperparams_path,
last_run_dir=raw.get("last_run_dir"),
)
@@ -153,24 +135,25 @@ def load_config(path: str | Path) -> AppConfig:
raw = yaml.safe_load(f)
if not isinstance(raw, dict):
raise ValueError("YAML 根节点必须是字典")
cfg = AppConfig.from_dict(raw)
sr = cfg.split_ratios
if len(sr) != 3:
raise ValueError("split_ratios 必须为长度为 3 的列表 [train, val, test]")
if abs(sum(sr) - 1.0) > 1e-6:
raise ValueError(f"split_ratios 之和必须为 1，当前为 {sum(sr)}")
if cfg.split_mode not in ("random", "grouped_stratified"):
raise ValueError("split_mode 必须为 random 或 grouped_stratified")
if cfg.split_stratify_target not in ("BW_3dB", "IL", "V_pi"):
raise ValueError("split_stratify_target 必须为 BW_3dB、IL 或 V_pi")
if cfg.split_stratify_bins < 2:
raise ValueError("split_stratify_bins 必须 >= 2")
if cfg.outlier_strategy not in ("none", "iqr", "zscore", "quantile_clip"):
raise ValueError(f"未知 outlier_strategy: {cfg.outlier_strategy}")
if cfg.outlier_apply_to not in ("targets", "all"):
raise ValueError("outlier_apply_to 必须为 targets 或 all")
if len(cfg.loss.target_weights) != 3:
raise ValueError("loss.target_weights 长度必须为 3")
if cfg.filter_v_pi_range and cfg.v_pi_min >= cfg.v_pi_max:
raise ValueError("启用 filter_v_pi_range 时须满足 v_pi_min < v_pi_max")
cfg = AppConfig.from_dict(raw, path.parent.resolve())
if not 0.0 < cfg.data.test_size < 1.0:
raise ValueError("data.test_size 必须在 (0, 1) 之间")
if cfg.data.filter_v_pi_max <= 0:
raise ValueError("data.filter_v_pi_max 必须 > 0")
if cfg.model.input_dim != 8:
raise ValueError("model.input_dim 必须为 8")
if cfg.model.output_dim != 3:
raise ValueError("model.output_dim 必须为 3")
if not cfg.model.hidden_dims:
raise ValueError("model.hidden_dims 不能为空")
if cfg.model.n_experts < 1:
raise ValueError("model.n_experts 必须 >= 1")
if cfg.model.gating_hidden < 1:
raise ValueError("model.gating_hidden 必须 >= 1")
if cfg.model.activation not in ("relu", "gaussian"):
raise ValueError("model.activation 必须为 relu 或 gaussian")
if len(cfg.optimizer.betas) != 2:
raise ValueError("optimizer.betas 长度必须为 2")
if cfg.training.batch_size < 1 or cfg.training.epochs < 1:
raise ValueError("training.batch_size 与 training.epochs 必须 >= 1")
return cfg

75
src/evaluate.py
View File

@@ -1,12 +1,12 @@
"""加载最优模型并在各划分上评估，导出 CSV / JSON。"""
"""加载模型并在 train/test 上评估，导出 CSV / JSON。"""
from __future__ import annotations
import csv
import json
import logging
from pathlib import Path
from dataclasses import asdict
from pathlib import Path
from typing import Dict, Tuple
import numpy as np
@@ -16,15 +16,10 @@ import torch.nn as nn
from src.config import AppConfig
from src.data import INPUT_COLUMNS, TARGET_COLUMNS
from src.losses import build_loss
from src.metrics import (
FullMetricsReport,
compute_full_report,
report_to_flat_dict,
)
from src.model import MLPRegressor
from src.preprocess import ProcessedDataBundle
from src.trainer import evaluate_loss_loader, load_weights
from src.metrics import FullMetricsReport, compute_full_report, report_to_flat_dict
from src.model import create_model_from_config
from src.preprocess import ProcessedDataBundle, inverse_transform_targets
from src.trainer import load_weights
logger = logging.getLogger(__name__)
@@ -40,25 +35,32 @@ def predict_all(
for xb, yb in loader:
xb = xb.to(device)
pr = model(xb).detach().cpu().numpy()
yt = yb.numpy()
preds.append(pr)
trues.append(yt)
trues.append(yb.numpy())
return np.concatenate(preds, axis=0), np.concatenate(trues, axis=0)
def _select_checkpoint(run_dir: Path) -> Path:
ckpt_last = run_dir / "checkpoints" / "last.pt"
if ckpt_last.is_file():
return ckpt_last
ckpt_best = run_dir / "checkpoints" / "best.pt"
if ckpt_best.is_file():
return ckpt_best
raise FileNotFoundError(f"未找到 {run_dir}/checkpoints/last.pt 或 best.pt")
def evaluate_split(
model: nn.Module,
criterion: nn.Module,
loader: torch.utils.data.DataLoader,
device: torch.device,
y_scaler,
y_scalers,
split_name: str,
) -> Tuple[FullMetricsReport, FullMetricsReport, float]:
"""返回 (标准化空间报告, 物理空间报告, 平均损失)。"""
loss = evaluate_loss_loader(model, loader, criterion, device)
pred_n, true_n = predict_all(model, loader, device)
pred_p = y_scaler.inverse_transform(pred_n)
true_p = y_scaler.inverse_transform(true_n)
loss = float(nn.functional.mse_loss(torch.from_numpy(pred_n), torch.from_numpy(true_n)).item())
pred_p = inverse_transform_targets(pred_n, y_scalers)
true_p = inverse_transform_targets(true_n, y_scalers)
rep_n = compute_full_report(split_name, loss, true_n, pred_n, TARGET_COLUMNS)
rep_p = compute_full_report(split_name, loss, true_p, pred_p, TARGET_COLUMNS)
return rep_n, rep_p, loss
@@ -70,31 +72,18 @@ def run_full_evaluation(
run_dir: Path,
device: torch.device,
) -> Tuple[nn.Module, Dict]:
"""载入 best.pt，在 train/val/test 上评估并写 metrics.csv 与 summary.json；返回模型与摘要。"""
model = MLPRegressor(
input_dim=cfg.model.input_dim,
hidden_dims=cfg.model.hidden_dims,
output_dim=cfg.model.output_dim,
batchnorm=cfg.model.batchnorm,
dropout=cfg.model.dropout,
residual=cfg.model.residual,
).to(device)
ckpt_best = run_dir / "checkpoints" / "best.pt"
load_weights(model, ckpt_best, device)
criterion = build_loss(cfg.loss).to(device)
y_scaler = bundle.y_scaler
model = create_model_from_config(cfg).to(device)
load_weights(model, _select_checkpoint(run_dir), device)
rows = []
summary: Dict = {"splits": {}}
for name, loader in (
("train", bundle.train_loader),
("val", bundle.val_loader),
("test", bundle.test_loader),
):
for name, loader in (("train", bundle.train_loader), ("test", bundle.test_loader)):
rep_n, rep_p, loss = evaluate_split(
model, criterion, loader, device, y_scaler, name
model,
loader,
device,
bundle.y_scalers,
name,
)
summary["splits"][name] = {
"loss": loss,
@@ -106,7 +95,6 @@ def run_full_evaluation(
rows.append(row)
metrics_path = run_dir / "metrics.csv"
if rows:
with metrics_path.open("w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
writer.writeheader()
@@ -125,11 +113,10 @@ def export_test_predictions_csv(
device: torch.device,
path: Path,
) -> None:
"""导出测试集物理空间真值、预测与误差。"""
model.eval()
pred_n, true_n = predict_all(model, bundle.test_loader, device)
pred_p = bundle.y_scaler.inverse_transform(pred_n)
true_p = bundle.y_scaler.inverse_transform(true_n)
pred_p = inverse_transform_targets(pred_n, bundle.y_scalers)
true_p = inverse_transform_targets(true_n, bundle.y_scalers)
err = pred_p - true_p
cols: Dict[str, np.ndarray] = {}
for j, name in enumerate(INPUT_COLUMNS):

23
src/infer.py
View File

@@ -6,13 +6,14 @@ import logging
from pathlib import Path
from typing import Optional
import numpy as np
import pandas as pd
import torch
from src.config import AppConfig, load_config
from src.data import INPUT_COLUMNS, TARGET_COLUMNS, _strip_optional_list_brackets
from src.model import MLPRegressor
from src.preprocess import load_scalers
from src.model import create_model_from_config
from src.preprocess import inverse_transform_targets, load_scalers
from src.trainer import load_weights
logger = logging.getLogger(__name__)
@@ -68,23 +69,19 @@ def run_inference(
载入 best 模型与 scaler，对输入表进行批量推理并写出 CSV（物理量空间）。
"""
X = _read_inputs_table(Path(input_path)).to_numpy(dtype=np.float32)
X_scaler, y_scaler = load_scalers(run_dir)
X_scaler, y_scalers = load_scalers(run_dir)
Xn = X_scaler.transform(X)
model = MLPRegressor(
input_dim=cfg.model.input_dim,
hidden_dims=cfg.model.hidden_dims,
output_dim=cfg.model.output_dim,
batchnorm=cfg.model.batchnorm,
dropout=cfg.model.dropout,
residual=cfg.model.residual,
).to(device)
load_weights(model, run_dir / "checkpoints" / "best.pt", device)
model = create_model_from_config(cfg).to(device)
ckpt = run_dir / "checkpoints" / "last.pt"
if not ckpt.is_file():
ckpt = run_dir / "checkpoints" / "best.pt"
load_weights(model, ckpt, device)
model.eval()
with torch.no_grad():
pred_n = model(torch.from_numpy(Xn).float().to(device)).cpu().numpy()
pred_p = y_scaler.inverse_transform(pred_n)
pred_p = inverse_transform_targets(pred_n, y_scalers)
out = pd.DataFrame(X, columns=INPUT_COLUMNS)
for j, name in enumerate(TARGET_COLUMNS):

26
src/main.py
View File

@@ -13,7 +13,7 @@ import torch
from src.config import load_config
from src.data import load_raw_txt, quality_report_before_clean, summarize_for_console
from src.evaluate import export_test_predictions_csv, run_full_evaluation
from src.model import MLPRegressor
from src.model import create_model_from_config, weights_init
from src.plots import generate_all_figures
from src.preprocess import prepare_training_data, rebuild_bundle_for_eval
from src.trainer import fit
@@ -61,7 +61,7 @@ def _write_summary_md(run_dir: Path, summary: dict) -> None:
lines.append(f"## {split}")
lines.append("")
lines.append(
f"- **损失（标准化输出空间 Huber/MSE 准则）**: {block['loss']:.6f}"
f"- **损失（标准化输出空间 MSE）**: {block['loss']:.6f}"
)
for space, label in ("normalized", "标准化空间"), ("physical", "物理量空间"):
sub = block[space]
@@ -78,7 +78,7 @@ def _write_summary_md(run_dir: Path, summary: dict) -> None:
def cmd_train(args: argparse.Namespace) -> None:
cfg_path = _resolve_cfg_path(args.config)
cfg = load_config(cfg_path)
set_global_seed(cfg.random_seed)
set_global_seed(cfg.data.random_state)
run_dir = make_run_dir(resolve_path(cfg.output_dir, _project_root()))
shutil.copy2(cfg_path, run_dir / "config_snapshot.yaml")
@@ -86,21 +86,15 @@ def cmd_train(args: argparse.Namespace) -> None:
data_path = resolve_path(cfg.data_path, _project_root())
df = load_raw_txt(data_path)
q = quality_report_before_clean(df, cfg.v_pi_min, cfg.v_pi_max)
q = quality_report_before_clean(df, 0.0, cfg.data.filter_v_pi_max)
logger.info("数据质量(清洗前): %s", summarize_for_console(df, q))
bundle = prepare_training_data(df, cfg, run_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MLPRegressor(
input_dim=cfg.model.input_dim,
hidden_dims=cfg.model.hidden_dims,
output_dim=cfg.model.output_dim,
batchnorm=cfg.model.batchnorm,
dropout=cfg.model.dropout,
residual=cfg.model.residual,
).to(device)
model = create_model_from_config(cfg).to(device)
model.apply(weights_init)
history = fit(model, cfg, bundle.train_loader, bundle.val_loader, run_dir, device)
history = fit(model, cfg, bundle, run_dir, device)
model_eval, summary = run_full_evaluation(cfg, bundle, run_dir, device)
export_test_predictions_csv(
bundle, model_eval, device, run_dir / "test_predictions.csv"
@@ -115,7 +109,7 @@ def cmd_eval(args: argparse.Namespace) -> None:
run_dir = _resolve_run_dir(cfg_path, args.run_dir, args.output_dir)
snap = run_dir / "config_snapshot.yaml"
cfg = load_config(snap if snap.is_file() else cfg_path)
set_global_seed(cfg.random_seed)
set_global_seed(cfg.data.random_state)
setup_logging(run_dir / "eval.log")
data_path = resolve_path(cfg.data_path, _project_root())
@@ -135,7 +129,7 @@ def cmd_infer(args: argparse.Namespace) -> None:
run_dir = _resolve_run_dir(cfg_path, args.run_dir, args.output_dir)
snap = run_dir / "config_snapshot.yaml"
cfg = load_config(snap if snap.is_file() else cfg_path)
set_global_seed(cfg.random_seed)
set_global_seed(cfg.data.random_state)
setup_logging(None)
from src.infer import run_inference
@@ -147,7 +141,7 @@ def cmd_infer(args: argparse.Namespace) -> None:
def build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(description="MZM MLP 训练 / 评估 / 推理")
p = argparse.ArgumentParser(description="MZM MoE PINN 训练 / 评估 / 推理")
sub = p.add_subparsers(dest="command", required=True)
pt = sub.add_parser("train", help="训练模型")

129
src/model.py
View File

@@ -1,4 +1,4 @@
"""可配置 MLP 回归模型。"""
"""MoE PINN 模型定义。"""
from __future__ import annotations
@@ -8,54 +8,103 @@ import torch
import torch.nn as nn
def kaiming_init_module(m: nn.Module) -> None:
"""对 Linear 使用 Kaiming uniform（ReLU），偏置置零。"""
if isinstance(m, nn.Linear):
nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.BatchNorm1d):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
class GaussianActivation(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.exp(-(x**2))
class MLPRegressor(nn.Module):
"""
多层感知机回归：输入 8 维，输出 3 维。
def weights_init(layer_in: nn.Module) -> None:
"""与 notebook 保持一致的 Kaiming 初始化。"""
if isinstance(layer_in, nn.Linear):
nn.init.kaiming_uniform_(layer_in.weight)
if layer_in.bias is not None:
layer_in.bias.data.fill_(0.0)
可选 BatchNorm1d、Dropout、以及在相邻层维度相等时的残差相加。
"""
def build_activation(name: str) -> nn.Module:
if name == "relu":
return nn.ReLU()
if name == "gaussian":
return GaussianActivation()
raise ValueError(f"未知激活函数: {name}")
class ExpertNN(nn.Module):
def __init__(
self,
input_dim: int,
hidden_dims: List[int],
output_dim: int,
batchnorm: bool = False,
dropout: float = 0.0,
residual: bool = False,
hidden_dims: List[int],
activation_fn: nn.Module,
dropout_rate: float = 0.0,
use_bn: bool = False,
) -> None:
super().__init__()
self.residual = residual
dims = [input_dim] + list(hidden_dims) + [output_dim]
self._hidden_blocks = nn.ModuleList()
for i in range(len(dims) - 2):
in_d, out_d = dims[i], dims[i + 1]
seq_layers: list[nn.Module] = [nn.Linear(in_d, out_d)]
if batchnorm:
seq_layers.append(nn.BatchNorm1d(out_d))
seq_layers.append(nn.ReLU(inplace=True))
if dropout and dropout > 0:
seq_layers.append(nn.Dropout(p=dropout))
self._hidden_blocks.append(nn.Sequential(*seq_layers))
self._head = nn.Linear(dims[-2], dims[-1])
self.apply(kaiming_init_module)
layers: list[nn.Module] = []
prev_dim = input_dim
for h in hidden_dims:
layers.append(nn.Linear(prev_dim, h))
if use_bn:
layers.append(nn.BatchNorm1d(h))
layers.append(type(activation_fn)() if isinstance(activation_fn, nn.ReLU) else activation_fn.__class__())
if dropout_rate > 0:
layers.append(nn.Dropout(p=dropout_rate))
prev_dim = h
layers.append(nn.Linear(prev_dim, output_dim))
self.net = nn.Sequential(*layers)
def forward(self, x: torch.Tensor) -> torch.Tensor:
h = x
for block in self._hidden_blocks:
inp = h
h = block(inp)
if self.residual and inp.shape[-1] == h.shape[-1]:
h = h + inp
return self._head(h)
return self.net(x)
class MixtureOfExperts(nn.Module):
def __init__(
self,
input_dim: int,
output_dim: int,
hidden_dims: List[int],
n_experts: int = 3,
activation_fn: nn.Module | None = None,
gating_hidden: int = 32,
dropout_rate: float = 0.0,
use_bn: bool = False,
) -> None:
super().__init__()
act = activation_fn if activation_fn is not None else nn.ReLU()
self.experts = nn.ModuleList(
[
ExpertNN(
input_dim,
output_dim,
hidden_dims,
act,
dropout_rate=dropout_rate,
use_bn=use_bn,
)
for _ in range(n_experts)
]
)
self.gating = nn.Sequential(
nn.Linear(input_dim, gating_hidden),
nn.ReLU(),
nn.Linear(gating_hidden, n_experts),
nn.Softmax(dim=1),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
gate_weights = self.gating(x)
expert_outputs = torch.stack([expert(x) for expert in self.experts], dim=2)
return torch.bmm(expert_outputs, gate_weights.unsqueeze(2)).squeeze(2)
def create_model_from_config(cfg) -> MixtureOfExperts:
return MixtureOfExperts(
input_dim=cfg.model.input_dim,
output_dim=cfg.model.output_dim,
hidden_dims=cfg.model.hidden_dims,
n_experts=cfg.model.n_experts,
activation_fn=build_activation(cfg.model.activation),
gating_hidden=cfg.model.gating_hidden,
dropout_rate=cfg.model.dropout_rate,
use_bn=cfg.model.use_bn,
)

28
src/plots.py
View File

@@ -13,22 +13,22 @@ import torch
from src.config import AppConfig
from src.data import TARGET_COLUMNS
from src.evaluate import predict_all
from src.model import MLPRegressor
from src.preprocess import ProcessedDataBundle
from src.model import create_model_from_config
from src.preprocess import ProcessedDataBundle, inverse_transform_targets
from src.trainer import TrainHistory, load_weights
logger = logging.getLogger(__name__)
def plot_loss_curves(history: TrainHistory, out_path: Path) -> None:
"""绘制 train/val loss 曲线。"""
"""绘制 train/test loss 曲线。"""
sns.set_theme(style="whitegrid", context="talk")
fig, ax = plt.subplots(figsize=(8, 5))
ax.plot(history.epoch, history.train_loss, label="Train loss", linewidth=2)
ax.plot(history.epoch, history.val_loss, label="Val loss", linewidth=2)
ax.plot(history.epoch, history.test_loss, label="Test loss", linewidth=2)
ax.set_xlabel("Epoch")
ax.set_ylabel("Loss (normalized target space)")
ax.set_title("Training / Validation Loss")
ax.set_title("Training / Test Loss")
ax.legend()
fig.tight_layout()
out_path.parent.mkdir(parents=True, exist_ok=True)
@@ -104,19 +104,15 @@ def generate_all_figures(
fig_dir.mkdir(parents=True, exist_ok=True)
plot_loss_curves(history, fig_dir / "loss_curve.png")
model = MLPRegressor(
input_dim=cfg.model.input_dim,
hidden_dims=cfg.model.hidden_dims,
output_dim=cfg.model.output_dim,
batchnorm=cfg.model.batchnorm,
dropout=cfg.model.dropout,
residual=cfg.model.residual,
).to(device)
load_weights(model, run_dir / "checkpoints" / "best.pt", device)
model = create_model_from_config(cfg).to(device)
ckpt = run_dir / "checkpoints" / "last.pt"
if not ckpt.is_file():
ckpt = run_dir / "checkpoints" / "best.pt"
load_weights(model, ckpt, device)
pred_n, true_n = predict_all(model, bundle.test_loader, device)
pred_p = bundle.y_scaler.inverse_transform(pred_n)
true_p = bundle.y_scaler.inverse_transform(true_n)
pred_p = inverse_transform_targets(pred_n, bundle.y_scalers)
true_p = inverse_transform_targets(true_n, bundle.y_scalers)
for i, name in enumerate(TARGET_COLUMNS):
plot_scatter_true_pred(

547
src/preprocess.py
View File

@@ -7,7 +7,7 @@ import logging
import pickle
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from typing import List, Sequence, Tuple
import numpy as np
import pandas as pd
@@ -17,331 +17,55 @@ from sklearn.preprocessing import StandardScaler
from torch.utils.data import DataLoader, TensorDataset
from src.config import AppConfig
from src.data import ALL_COLUMNS, INPUT_COLUMNS, TARGET_COLUMNS, quality_report_before_clean
from src.data import INPUT_COLUMNS, TARGET_COLUMNS, quality_report_before_clean
logger = logging.getLogger(__name__)
@dataclass
class ProcessedDataBundle:
"""训练用张量与 DataLoader，以及划分后的 numpy（含测试集原始物理量用于导出）。"""
"""训练用张量与 DataLoader，以及 notebook 风格的 train/test 划分数据。"""
train_loader: DataLoader
val_loader: DataLoader
test_loader: DataLoader
X_train: np.ndarray
X_val: np.ndarray
X_test: np.ndarray
y_train: np.ndarray
y_val: np.ndarray
y_test: np.ndarray
X_train_raw: np.ndarray
X_test_raw: np.ndarray
y_train_raw: np.ndarray
y_test_raw: np.ndarray
X_scaler: StandardScaler
y_scaler: StandardScaler
y_scalers: List[StandardScaler]
feature_names: List[str]
target_names: List[str]
def _mask_outliers_iqr(
values: np.ndarray, col_names: List[str], k: float
) -> np.ndarray:
"""返回 True 表示该行在任一选定列上超出训练集 IQR 范围（基于传入的 values 统计）。"""
mask = np.zeros(len(values), dtype=bool)
for j, _ in enumerate(col_names):
col = values[:, j]
q1, q3 = np.percentile(col, [25, 75])
iqr = q3 - q1
lo, hi = q1 - k * iqr, q3 + k * iqr
mask |= (col < lo) | (col > hi)
return mask
def _mask_outliers_zscore(values: np.ndarray, threshold: float) -> np.ndarray:
mask = np.zeros(len(values), dtype=bool)
for j in range(values.shape[1]):
col = values[:, j]
mu, sig = col.mean(), col.std(ddof=0)
if sig < 1e-12:
continue
z = np.abs((col - mu) / sig)
mask |= z > threshold
return mask
def _winsorize_train_apply_all(
train: np.ndarray,
val: np.ndarray,
test: np.ndarray,
ql: float,
qu: float,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""按训练集分位数对 train/val/test 同步裁剪（列方向）。"""
lo = np.quantile(train, ql, axis=0)
hi = np.quantile(train, qu, axis=0)
def clip_arr(a: np.ndarray) -> np.ndarray:
return np.clip(a, lo, hi)
return clip_arr(train), clip_arr(val), clip_arr(test)
def clean_dataframe(
df: pd.DataFrame,
cfg: AppConfig,
report_lines: List[str],
) -> pd.DataFrame:
"""
清洗流程（顺序固定，便于复现与审计）：
1. 可选：完全重复行去重。
2. 可选：以 **最后一列对应字段 V_pi**（txt 第 11 个逗号分隔字段）为门控，仅保留
``v_pi_min <= V_pi <= v_pi_max``（默认 [0, 500]）。
3. 可选：再移除 ``V_pi <= 0``（与区间门控独立，由配置控制）。
"""
"""按 notebook 逻辑清洗：仅保留 V_pi < 阈值。"""
out = df.copy()
n0 = len(out)
if cfg.remove_duplicate_rows:
out = out.drop_duplicates()
report_lines.append(f"去完全重复行: {n0} -> {len(out)}")
if cfg.filter_v_pi_range:
n1 = len(out)
lo, hi = float(cfg.v_pi_min), float(cfg.v_pi_max)
mask = (out["V_pi"] >= lo) & (out["V_pi"] <= hi)
out = out[mask].reset_index(drop=True)
report_lines.append(
f"V_pi 物理区间过滤 [{lo}, {hi}]（txt 第 11 列 / 列名 V_pi）: {n1} -> {len(out)}"
)
if cfg.remove_nonpositive_vpi:
n2 = len(out)
out = out[out["V_pi"] > 0].reset_index(drop=True)
report_lines.append(f"移除 V_pi<=0: {n2} -> {len(out)}")
vmax = float(cfg.data.filter_v_pi_max)
out = out[out["V_pi"] < vmax].reset_index(drop=True)
report_lines.append(f"V_pi 阈值过滤 (< {vmax})：{n0} -> {len(out)}")
if len(out) == 0:
raise ValueError(
"清洗后样本数为 0：请检查 V_pi 区间配置、数据源或是否过度去重。"
)
raise ValueError("清洗后样本数为 0，请检查数据源或 V_pi 阈值设置。")
return out
def _random_split_indices(
n: int,
ratios: List[float],
seed: int,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""返回 train/val/test 的整数索引（先 shuffle 再按比例切分）。"""
rng = np.random.default_rng(seed)
idx = np.arange(n)
rng.shuffle(idx)
tr, va, te = ratios
n_test = int(round(n * te))
n_val = int(round(n * va))
n_train = n - n_val - n_test
if n_train <= 0 or n_val <= 0 or n_test <= 0:
raise ValueError(
f"划分后样本过少: train={n_train}, val={n_val}, test={n_test}，请调整比例或数据量"
)
i_train = idx[:n_train]
i_val = idx[n_train : n_train + n_val]
i_test = idx[n_train + n_val :]
return i_train, i_val, i_test
def _quantile_bin_labels(values: np.ndarray, n_bins: int) -> np.ndarray | None:
"""
基于秩做近似等频分桶，避免重复值导致的 qcut 退化。
返回每个样本所属桶标签；若样本过少则返回 None。
"""
if len(values) < 2:
return None
q = min(int(n_bins), len(values))
if q < 2:
return None
ranks = pd.Series(values).rank(method="first")
labels = pd.qcut(ranks, q=q, labels=False, duplicates="drop")
if labels is None:
return None
arr = np.asarray(labels, dtype=int)
if len(np.unique(arr)) < 2:
return None
return arr
def _grouped_split_indices(
df: pd.DataFrame,
ratios: List[float],
seed: int,
stratify_target: str,
stratify_bins: int,
report_lines: List[str],
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
先按输入 8 维分组，再在组级别按目标统计量近似分层切分。
这样可避免“同输入异输出”同时落在 train/val/test，提升评估稳定性。
若分层条件不足，则退化为组级随机切分。
"""
if len(df) < 3:
raise ValueError("样本数过少，无法做 train/val/test 切分。")
group_ids = df.groupby(INPUT_COLUMNS, sort=False, dropna=False).ngroup().to_numpy()
n_groups = int(group_ids.max()) + 1
group_df = df.copy()
group_df["_group_id"] = group_ids
group_stat = (
group_df.groupby("_group_id", sort=True)
.agg(group_size=("V_pi", "size"), strat_value=(stratify_target, "median"))
.reset_index()
)
group_id_arr = group_stat["_group_id"].to_numpy(dtype=int)
labels = _quantile_bin_labels(
group_stat["strat_value"].to_numpy(dtype=np.float64),
stratify_bins,
)
tr, va, te = ratios
holdout_ratio = va + te
val_ratio_in_holdout = va / holdout_ratio
def _split_groups(use_stratify: bool) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
strat = labels if use_stratify and labels is not None else None
train_groups, holdout_groups = train_test_split(
group_id_arr,
train_size=tr,
test_size=holdout_ratio,
random_state=seed,
shuffle=True,
stratify=strat,
)
holdout_strat = None
if strat is not None:
label_map = dict(zip(group_id_arr.tolist(), labels.tolist()))
holdout_labels = np.asarray(
[label_map[int(g)] for g in holdout_groups],
dtype=int,
)
if len(np.unique(holdout_labels)) >= 2:
holdout_strat = holdout_labels
val_groups, test_groups = train_test_split(
holdout_groups,
train_size=val_ratio_in_holdout,
test_size=1.0 - val_ratio_in_holdout,
random_state=seed + 1,
shuffle=True,
stratify=holdout_strat,
)
return (
np.asarray(train_groups, dtype=int),
np.asarray(val_groups, dtype=int),
np.asarray(test_groups, dtype=int),
)
split_note = (
f"按输入分组切分，共 {n_groups} 个唯一输入组；"
f"组级按 {stratify_target} 中位数分 {min(stratify_bins, n_groups)} 桶近似分层。"
)
try:
train_groups, val_groups, test_groups = _split_groups(use_stratify=True)
report_lines.append(split_note)
except ValueError as e:
train_groups, val_groups, test_groups = _split_groups(use_stratify=False)
report_lines.append(f"{split_note} 但分层条件不足，退化为组级随机切分：{e}")
i_train = np.flatnonzero(np.isin(group_ids, train_groups))
i_val = np.flatnonzero(np.isin(group_ids, val_groups))
i_test = np.flatnonzero(np.isin(group_ids, test_groups))
return i_train, i_val, i_test
def build_split_indices(
df: pd.DataFrame,
cfg: AppConfig,
report_lines: List[str],
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""根据配置生成 train/val/test 行索引。"""
if cfg.split_mode == "random":
report_lines.append("切分策略：随机打乱后按比例切分。")
return _random_split_indices(len(df), cfg.split_ratios, cfg.random_seed)
return _grouped_split_indices(
df=df,
ratios=cfg.split_ratios,
seed=cfg.random_seed,
stratify_target=cfg.split_stratify_target,
stratify_bins=cfg.split_stratify_bins,
report_lines=report_lines,
)
def apply_train_only_outliers(
X_train: np.ndarray,
y_train: np.ndarray,
X_val: np.ndarray,
y_val: np.ndarray,
X_test: np.ndarray,
y_test: np.ndarray,
cfg: AppConfig,
report_lines: List[str],
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""
仅在训练集上估计阈值：
- iqr/zscore: 从训练集删除离群行（val/test 不动）
- quantile_clip: 对 train/val/test 同步 winsorize（阈值来自 train）
"""
strat = cfg.outlier_strategy
if strat == "none":
report_lines.append("outlier_strategy=none：不对数值做裁剪/删除（除配置项外）。")
return X_train, y_train, X_val, y_val, X_test, y_test
cols = cfg.outlier_apply_to
if cols == "all":
train_mat = np.hstack([X_train, y_train])
val_mat = np.hstack([X_val, y_val])
test_mat = np.hstack([X_test, y_test])
names = INPUT_COLUMNS + TARGET_COLUMNS
else:
train_mat = y_train.copy()
val_mat = y_val.copy()
test_mat = y_test.copy()
names = TARGET_COLUMNS
if strat == "quantile_clip":
ql = cfg.outlier_config.quantile_lower
qu = cfg.outlier_config.quantile_upper
tr2, va2, te2 = _winsorize_train_apply_all(train_mat, val_mat, test_mat, ql, qu)
report_lines.append(
f"quantile_clip: 按训练集分位数 [{ql}, {qu}] 对 {cols} 列 winsorize。"
)
if cols == "all":
d = len(INPUT_COLUMNS)
X_train, y_train = tr2[:, :d], tr2[:, d:]
X_val, y_val = va2[:, :d], va2[:, d:]
X_test, y_test = te2[:, :d], te2[:, d:]
else:
y_train, y_val, y_test = tr2, va2, te2
return X_train, y_train, X_val, y_val, X_test, y_test
if strat == "iqr":
mask = _mask_outliers_iqr(train_mat, names, cfg.outlier_config.iqr_k)
elif strat == "zscore":
mask = _mask_outliers_zscore(train_mat, cfg.outlier_config.zscore_threshold)
else:
raise ValueError(f"未知 outlier_strategy: {strat}")
removed = int(mask.sum())
kept = ~mask
X_train, y_train = X_train[kept], y_train[kept]
report_lines.append(
f"{strat}: 在训练子集上检测 {cols} 离群，删除训练行 {removed}，保留 {len(X_train)}。"
)
return X_train, y_train, X_val, y_val, X_test, y_test
def build_dataloaders(
X_train: np.ndarray,
y_train: np.ndarray,
X_val: np.ndarray,
y_val: np.ndarray,
X_test: np.ndarray,
y_test: np.ndarray,
batch_size: int,
num_workers: int,
) -> Tuple[DataLoader, DataLoader, DataLoader]:
) -> Tuple[DataLoader, DataLoader]:
def to_loader(X: np.ndarray, y: np.ndarray, shuffle: bool) -> DataLoader:
ds = TensorDataset(
torch.from_numpy(X).float(),
@@ -357,127 +81,77 @@ def build_dataloaders(
return (
to_loader(X_train, y_train, shuffle=True),
to_loader(X_val, y_val, shuffle=False),
to_loader(X_test, y_test, shuffle=False),
)
def _fit_target_scalers(y_train: np.ndarray) -> Tuple[List[StandardScaler], np.ndarray]:
scalers: List[StandardScaler] = []
scaled_cols = []
for i in range(y_train.shape[1]):
scaler = StandardScaler()
scaled_cols.append(scaler.fit_transform(y_train[:, i : i + 1]))
scalers.append(scaler)
return scalers, np.hstack(scaled_cols)
def transform_targets(y: np.ndarray, y_scalers: Sequence[StandardScaler]) -> np.ndarray:
cols = [scaler.transform(y[:, i : i + 1]) for i, scaler in enumerate(y_scalers)]
return np.hstack(cols)
def inverse_transform_targets(
y_scaled: np.ndarray, y_scalers: Sequence[StandardScaler]
) -> np.ndarray:
cols = [scaler.inverse_transform(y_scaled[:, i : i + 1]) for i, scaler in enumerate(y_scalers)]
return np.hstack(cols)
def save_scalers(
X_scaler: StandardScaler,
y_scaler: StandardScaler,
y_scalers: Sequence[StandardScaler],
run_dir: Path,
) -> None:
with (run_dir / "x_scaler.pkl").open("wb") as f:
pickle.dump(X_scaler, f)
with (run_dir / "y_scaler.pkl").open("wb") as f:
pickle.dump(y_scaler, f)
with (run_dir / "y_scalers.pkl").open("wb") as f:
pickle.dump(list(y_scalers), f)
def load_scalers(run_dir: Path) -> Tuple[StandardScaler, StandardScaler]:
def load_scalers(run_dir: Path) -> Tuple[StandardScaler, List[StandardScaler]]:
with (run_dir / "x_scaler.pkl").open("rb") as f:
X_scaler = pickle.load(f)
with (run_dir / "y_scaler.pkl").open("rb") as f:
y_scaler = pickle.load(f)
return X_scaler, y_scaler
with (run_dir / "y_scalers.pkl").open("rb") as f:
y_scalers = pickle.load(f)
return X_scaler, list(y_scalers)
def save_split_indices(
run_dir: Path,
i_train: np.ndarray,
i_val: np.ndarray,
i_test: np.ndarray,
) -> None:
"""保存对清洗后矩阵行的划分索引，便于 eval 阶段完全复现。"""
payload = {
"train": i_train.astype(int).tolist(),
"val": i_val.astype(int).tolist(),
"test": i_test.astype(int).tolist(),
}
(run_dir / "split_indices.json").write_text(
json.dumps(payload, indent=2), encoding="utf-8"
json.dumps(payload, indent=2),
encoding="utf-8",
)
def load_split_indices(run_dir: Path) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
def load_split_indices(run_dir: Path) -> Tuple[np.ndarray, np.ndarray]:
path = run_dir / "split_indices.json"
if not path.is_file():
raise FileNotFoundError(
f"未找到 {path}。请使用本仓库训练产生的 run 目录，或先完成一次训练。"
)
raise FileNotFoundError(f"未找到 {path}")
data = json.loads(path.read_text(encoding="utf-8"))
return (
np.asarray(data["train"], dtype=int),
np.asarray(data["val"], dtype=int),
np.asarray(data["test"], dtype=int),
)
def rebuild_bundle_for_eval(
df: pd.DataFrame,
cfg: AppConfig,
run_dir: Path,
) -> ProcessedDataBundle:
"""
与训练阶段相同的清洗、划分与离群处理，但使用已保存的 StandardScaler 仅做 transform。
用于独立 eval / infer 流程，避免重新拟合 scaler 造成分布偏移。
"""
report_lines: List[str] = []
cleaned = clean_dataframe(df, cfg, report_lines)
X = cleaned[INPUT_COLUMNS].to_numpy(dtype=np.float64)
y = cleaned[TARGET_COLUMNS].to_numpy(dtype=np.float64)
i_tr, i_va, i_te = load_split_indices(run_dir)
for name, idx in ("train", i_tr), ("val", i_va), ("test", i_te):
if len(idx) == 0 or int(idx.max()) >= len(X) or int(idx.min()) < 0:
raise ValueError(
f"split_indices.json 与当前数据不兼容（{name} 索引越界或为空）。"
f"请确认 data_path 指向与训练相同的清洗后样本空间。"
)
X_train, y_train = X[i_tr], y[i_tr]
X_val, y_val = X[i_va], y[i_va]
X_test, y_test = X[i_te], y[i_te]
X_train, y_train, X_val, y_val, X_test, y_test = apply_train_only_outliers(
X_train, y_train, X_val, y_val, X_test, y_test, cfg, report_lines
)
X_scaler, y_scaler = load_scalers(run_dir)
X_train_s = X_scaler.transform(X_train)
y_train_s = y_scaler.transform(y_train)
X_val_s = X_scaler.transform(X_val)
y_val_s = y_scaler.transform(y_val)
X_test_s = X_scaler.transform(X_test)
y_test_s = y_scaler.transform(y_test)
train_loader, val_loader, test_loader = build_dataloaders(
X_train_s,
y_train_s,
X_val_s,
y_val_s,
X_test_s,
y_test_s,
cfg.training.batch_size,
cfg.training.num_workers,
)
return ProcessedDataBundle(
train_loader=train_loader,
val_loader=val_loader,
test_loader=test_loader,
X_train=X_train_s,
X_val=X_val_s,
X_test=X_test_s,
y_train=y_train_s,
y_val=y_val_s,
y_test=y_test_s,
X_test_raw=X_test,
y_test_raw=y_test,
X_scaler=X_scaler,
y_scaler=y_scaler,
feature_names=list(INPUT_COLUMNS),
target_names=list(TARGET_COLUMNS),
)
def write_data_report_md(
path: Path,
raw_quality: dict,
@@ -515,12 +189,9 @@ def prepare_training_data(
cfg: AppConfig,
run_dir: Path,
) -> ProcessedDataBundle:
"""
完整预处理流水线：质量报告 -> 清洗 -> 划分 -> 训练集离群处理 -> 标准化 -> DataLoader。
将 data_report.md 与 cleaning 元数据写入 run_dir。
"""
"""按 notebook 一致逻辑准备 train/test、scaler 与 DataLoader。"""
report_lines: List[str] = []
raw_q = quality_report_before_clean(df, cfg.v_pi_min, cfg.v_pi_max)
raw_q = quality_report_before_clean(df, 0.0, cfg.data.filter_v_pi_max)
stats_before = df.describe().T
cleaned = clean_dataframe(df, cfg, report_lines)
@@ -528,39 +199,35 @@ def prepare_training_data(
X = cleaned[INPUT_COLUMNS].to_numpy(dtype=np.float64)
y = cleaned[TARGET_COLUMNS].to_numpy(dtype=np.float64)
all_idx = np.arange(len(X))
i_tr, i_va, i_te = build_split_indices(cleaned, cfg, report_lines)
save_split_indices(run_dir, i_tr, i_va, i_te)
X_train, y_train = X[i_tr], y[i_tr]
X_val, y_val = X[i_va], y[i_va]
X_test, y_test = X[i_te], y[i_te]
X_train_raw, X_test_raw, y_train_raw, y_test_raw, i_train, i_test = train_test_split(
X,
y,
all_idx,
test_size=cfg.data.test_size,
random_state=cfg.data.random_state,
)
report_lines.append(
f"划分 train/val/test = {cfg.split_ratios}，样本数 "
f"{len(X_train)}/{len(X_val)}/{len(X_test)}"
)
X_train, y_train, X_val, y_val, X_test, y_test = apply_train_only_outliers(
X_train, y_train, X_val, y_val, X_test, y_test, cfg, report_lines
f"train_test_split(test_size={cfg.data.test_size}, random_state={cfg.data.random_state}) "
f"-> {len(X_train_raw)}/{len(X_test_raw)}"
)
save_split_indices(run_dir, i_train, i_test)
X_scaler = StandardScaler()
y_scaler = StandardScaler()
X_train_s = X_scaler.fit_transform(X_train)
y_train_s = y_scaler.fit_transform(y_train)
X_val_s = X_scaler.transform(X_val)
y_val_s = y_scaler.transform(y_val)
X_test_s = X_scaler.transform(X_test)
y_test_s = y_scaler.transform(y_test)
save_scalers(X_scaler, y_scaler, run_dir)
X_train = X_scaler.fit_transform(X_train_raw)
X_test = X_scaler.transform(X_test_raw)
y_scalers, y_train = _fit_target_scalers(y_train_raw)
y_test = transform_targets(y_test_raw, y_scalers)
save_scalers(X_scaler, y_scalers, run_dir)
meta = {
"raw_quality": raw_q,
"cleaning_steps": report_lines,
"split_ratios": cfg.split_ratios,
"n_train": int(len(X_train_s)),
"n_val": int(len(X_val_s)),
"n_test": int(len(X_test_s)),
"test_size": cfg.data.test_size,
"random_state": cfg.data.random_state,
"n_train": int(len(X_train)),
"n_test": int(len(X_test)),
}
(run_dir / "cleaning_meta.json").write_text(
json.dumps(meta, indent=2, ensure_ascii=False, default=str),
@@ -576,31 +243,77 @@ def prepare_training_data(
stats_after.to_csv(run_dir / "data_stats.csv", encoding="utf-8")
logger.info("预处理完成：%s", run_dir / "data_report.md")
train_loader, val_loader, test_loader = build_dataloaders(
X_train_s,
y_train_s,
X_val_s,
y_val_s,
X_test_s,
y_test_s,
train_loader, test_loader = build_dataloaders(
X_train,
y_train,
X_test,
y_test,
cfg.training.batch_size,
cfg.training.num_workers,
)
return ProcessedDataBundle(
train_loader=train_loader,
val_loader=val_loader,
test_loader=test_loader,
X_train=X_train_s,
X_val=X_val_s,
X_test=X_test_s,
y_train=y_train_s,
y_val=y_val_s,
y_test=y_test_s,
X_test_raw=X_test,
y_test_raw=y_test,
X_train=X_train,
X_test=X_test,
y_train=y_train,
y_test=y_test,
X_train_raw=X_train_raw,
X_test_raw=X_test_raw,
y_train_raw=y_train_raw,
y_test_raw=y_test_raw,
X_scaler=X_scaler,
y_scaler=y_scaler,
y_scalers=y_scalers,
feature_names=list(INPUT_COLUMNS),
target_names=list(TARGET_COLUMNS),
)
def rebuild_bundle_for_eval(
df: pd.DataFrame,
cfg: AppConfig,
run_dir: Path,
) -> ProcessedDataBundle:
"""使用训练时保存的切分索引与 scaler 重新构造 train/test 数据。"""
report_lines: List[str] = []
cleaned = clean_dataframe(df, cfg, report_lines)
X = cleaned[INPUT_COLUMNS].to_numpy(dtype=np.float64)
y = cleaned[TARGET_COLUMNS].to_numpy(dtype=np.float64)
i_train, i_test = load_split_indices(run_dir)
for name, idx in (("train", i_train), ("test", i_test)):
if len(idx) == 0 or int(idx.max()) >= len(X) or int(idx.min()) < 0:
raise ValueError(f"split_indices.json 与当前数据不兼容（{name} 索引越界或为空）")
X_train_raw, X_test_raw = X[i_train], X[i_test]
y_train_raw, y_test_raw = y[i_train], y[i_test]
X_scaler, y_scalers = load_scalers(run_dir)
X_train = X_scaler.transform(X_train_raw)
X_test = X_scaler.transform(X_test_raw)
y_train = transform_targets(y_train_raw, y_scalers)
y_test = transform_targets(y_test_raw, y_scalers)
train_loader, test_loader = build_dataloaders(
X_train,
y_train,
X_test,
y_test,
cfg.training.batch_size,
cfg.training.num_workers,
)
return ProcessedDataBundle(
train_loader=train_loader,
test_loader=test_loader,
X_train=X_train,
X_test=X_test,
y_train=y_train,
y_test=y_test,
X_train_raw=X_train_raw,
X_test_raw=X_test_raw,
y_train_raw=y_train_raw,
y_test_raw=y_test_raw,
X_scaler=X_scaler,
y_scalers=y_scalers,
feature_names=list(INPUT_COLUMNS),
target_names=list(TARGET_COLUMNS),
)

256
src/trainer.py
View File

@@ -1,4 +1,4 @@
"""训练循环、早停、调度器与 checkpoint。"""
"""Notebook 风格的 MoE + autograd PINN 训练循环。"""
from __future__ import annotations
@@ -6,17 +6,15 @@ import csv
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from typing import Dict, List, Tuple
import numpy as np
import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau
from tqdm import tqdm
from src.config import AppConfig
from src.losses import build_loss
from src.model import MLPRegressor
from src.preprocess import ProcessedDataBundle
logger = logging.getLogger(__name__)
@@ -25,8 +23,7 @@ logger = logging.getLogger(__name__)
class TrainHistory:
epoch: List[int]
train_loss: List[float]
val_loss: List[float]
lr: List[float]
test_loss: List[float]
def _move_batch(
@@ -36,145 +33,172 @@ def _move_batch(
return x.to(device), y.to(device)
def train_one_epoch(
model: nn.Module,
loader: torch.utils.data.DataLoader,
criterion: nn.Module,
optimizer: torch.optim.Optimizer,
device: torch.device,
) -> float:
model.train()
total, n = 0.0, 0
for batch in loader:
xb, yb = _move_batch(batch, device)
optimizer.zero_grad(set_to_none=True)
pred = model(xb)
loss = criterion(pred, yb)
loss.backward()
optimizer.step()
total += float(loss.detach().cpu()) * xb.size(0)
n += xb.size(0)
return total / max(n, 1)
@torch.no_grad()
def evaluate_loss_loader(
model: nn.Module,
loader: torch.utils.data.DataLoader,
criterion: nn.Module,
device: torch.device,
) -> float:
model.eval()
total, n = 0.0, 0
for batch in loader:
xb, yb = _move_batch(batch, device)
pred = model(xb)
loss = criterion(pred, yb)
total += float(loss.detach().cpu()) * xb.size(0)
n += xb.size(0)
return total / max(n, 1)
def build_optimizer_and_scheduler(
model: nn.Module, cfg: AppConfig
) -> Tuple[AdamW, object]:
opt = AdamW(
def build_optimizer(model: nn.Module, cfg: AppConfig) -> AdamW:
beta1, beta2 = cfg.optimizer.betas
return AdamW(
model.parameters(),
lr=cfg.optimizer.lr,
weight_decay=cfg.optimizer.weight_decay,
betas=(beta1, beta2),
)
if cfg.scheduler.type == "cosine":
sched: torch.optim.lr_scheduler._LRScheduler = CosineAnnealingLR(
opt, T_max=cfg.training.epochs, eta_min=cfg.scheduler.plateau_min_lr
)
elif cfg.scheduler.type == "plateau":
sched = ReduceLROnPlateau(
opt,
mode="min",
factor=cfg.scheduler.plateau_factor,
patience=cfg.scheduler.plateau_patience,
min_lr=cfg.scheduler.plateau_min_lr,
)
def compute_pinn_loss(
model: nn.Module,
x_batch: torch.Tensor,
y_batch: torch.Tensor,
criterion: nn.Module,
cfg: AppConfig,
) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
x_in = x_batch.detach().clone().requires_grad_(True)
pred = model(x_in)
data_loss = criterion(pred, y_batch)
bw_pred = pred[:, 0]
il_pred = pred[:, 1]
vpi_pred = pred[:, 2]
length_batch = x_in[:, 7]
losses: Dict[str, torch.Tensor] = {}
total = data_loss
if cfg.physics.lambda_bw_mon != 0:
grads = torch.autograd.grad(
bw_pred,
x_in,
grad_outputs=torch.ones_like(bw_pred),
create_graph=True,
)[0]
d_bw_d_l = grads[:, 7]
losses["bw_mon"] = torch.mean(torch.relu(d_bw_d_l) ** 2)
total = total + cfg.physics.lambda_bw_mon * losses["bw_mon"]
if cfg.physics.lambda_IL_mon != 0:
grads = torch.autograd.grad(
il_pred,
x_in,
grad_outputs=torch.ones_like(il_pred),
create_graph=True,
)[0]
d_il_d_l = grads[:, 7]
losses["IL_mon"] = torch.mean(torch.relu(-d_il_d_l) ** 2)
total = total + cfg.physics.lambda_IL_mon * losses["IL_mon"]
if cfg.physics.lambda_vpiL != 0:
vpi_l = vpi_pred * length_batch
grads = torch.autograd.grad(
vpi_l,
x_in,
grad_outputs=torch.ones_like(vpi_l),
create_graph=True,
)[0]
d_vpi_l_d_l = grads[:, 7]
losses["vpiL"] = torch.mean(d_vpi_l_d_l**2)
total = total + cfg.physics.lambda_vpiL * losses["vpiL"]
if cfg.physics.lambda_smooth != 0:
grads1 = torch.autograd.grad(
bw_pred,
x_in,
grad_outputs=torch.ones_like(bw_pred),
create_graph=True,
)[0]
d_bw_d_l = grads1[:, 7]
grads2 = torch.autograd.grad(
d_bw_d_l,
x_in,
grad_outputs=torch.ones_like(d_bw_d_l),
create_graph=True,
)[0]
d2_bw_d_l2 = grads2[:, 7]
losses["smooth"] = torch.mean(d2_bw_d_l2**2)
total = total + cfg.physics.lambda_smooth * losses["smooth"]
return total, data_loss, losses
@torch.no_grad()
def evaluate_full_batch_mse(
model: nn.Module,
x: np.ndarray | torch.Tensor,
y: np.ndarray | torch.Tensor,
device: torch.device,
) -> float:
model.eval()
if isinstance(x, torch.Tensor):
x_t = x.to(device)
else:
raise ValueError(f"未知 scheduler.type: {cfg.scheduler.type}")
return opt, sched
x_t = torch.from_numpy(x).float().to(device)
if isinstance(y, torch.Tensor):
y_t = y.to(device)
else:
y_t = torch.from_numpy(y).float().to(device)
pred = model(x_t)
return float(nn.functional.mse_loss(pred, y_t).item())
def fit(
model: nn.Module,
cfg: AppConfig,
train_loader: torch.utils.data.DataLoader,
val_loader: torch.utils.data.DataLoader,
bundle: ProcessedDataBundle,
run_dir: Path,
device: torch.device,
) -> TrainHistory:
"""
训练模型：早停依据验证集损失；保存 best / last 权重到 run_dir/checkpoints。
同步写入 train_log.csv。
"""
criterion = build_loss(cfg.loss).to(device)
optimizer, scheduler = build_optimizer_and_scheduler(model, cfg)
"""按 notebook 风格训练，并记录每轮全量 train/test MSE。"""
criterion = nn.MSELoss().to(device)
optimizer = build_optimizer(model, cfg)
ckpt_dir = run_dir / "checkpoints"
ckpt_dir.mkdir(parents=True, exist_ok=True)
log_path = run_dir / "train_log.csv"
best_val = float("inf")
best_epoch = -1
patience_left = cfg.training.early_stopping_patience
x_train_full = torch.from_numpy(bundle.X_train).float().to(device)
y_train_full = torch.from_numpy(bundle.y_train).float().to(device)
x_test_full = torch.from_numpy(bundle.X_test).float().to(device)
y_test_full = torch.from_numpy(bundle.y_test).float().to(device)
hist = TrainHistory(epoch=[], train_loss=[], val_loss=[], lr=[])
hist = TrainHistory(epoch=[], train_loss=[], test_loss=[])
best_test = float("inf")
with log_path.open("w", newline="", encoding="utf-8") as fcsv:
writer = csv.writer(fcsv)
writer.writerow(["epoch", "train_loss", "val_loss", "lr", "best_val"])
writer.writerow(["epoch", "train_loss", "test_loss"])
for epoch in range(1, cfg.training.epochs + 1):
tr_loss = train_one_epoch(model, train_loader, criterion, optimizer, device)
va_loss = evaluate_loss_loader(model, val_loader, criterion, device)
for epoch in range(cfg.training.epochs):
model.train()
for x_batch, y_batch in bundle.train_loader:
x_batch, y_batch = _move_batch((x_batch, y_batch), device)
optimizer.zero_grad()
loss, _, _ = compute_pinn_loss(model, x_batch, y_batch, criterion, cfg)
loss.backward()
optimizer.step()
if cfg.scheduler.type == "cosine":
scheduler.step()
elif cfg.scheduler.type == "plateau":
scheduler.step(va_loss)
train_loss = evaluate_full_batch_mse(model, x_train_full, y_train_full, device)
test_loss = evaluate_full_batch_mse(model, x_test_full, y_test_full, device)
lr_now = float(optimizer.param_groups[0]["lr"])
hist.epoch.append(epoch)
hist.train_loss.append(tr_loss)
hist.val_loss.append(va_loss)
hist.lr.append(lr_now)
improved = va_loss + 1e-12 < best_val
if improved:
best_val = va_loss
best_epoch = epoch
patience_left = cfg.training.early_stopping_patience
torch.save(
{"epoch": epoch, "model_state": model.state_dict(), "val_loss": va_loss},
ckpt_dir / "best.pt",
)
else:
patience_left -= 1
writer.writerow([epoch, tr_loss, va_loss, lr_now, best_val])
hist.train_loss.append(train_loss)
hist.test_loss.append(test_loss)
writer.writerow([epoch, train_loss, test_loss])
fcsv.flush()
if epoch % 10 == 0 or epoch == 0:
logger.info(
"Epoch %d | train_loss=%.6f val_loss=%.6f | best_val=%.6f @%d",
"Epoch %5d | Train %.6f | Test %.6f",
epoch,
tr_loss,
va_loss,
best_val,
best_epoch,
train_loss,
test_loss,
)
torch.save(
{"epoch": epoch, "model_state": model.state_dict(), "val_loss": va_loss},
ckpt_dir / "last.pt",
)
if patience_left <= 0:
logger.info("早停触发于 epoch %d，最佳 epoch=%d", epoch, best_epoch)
break
payload = {
"epoch": epoch,
"model_state": model.state_dict(),
"train_loss": train_loss,
"test_loss": test_loss,
}
torch.save(payload, ckpt_dir / "last.pt")
if test_loss < best_test:
best_test = test_loss
torch.save(payload, ckpt_dir / "best.pt")
return hist

234
tests/test_smoke.py
View File

@@ -1,4 +1,4 @@
"""最小冒烟测试：模型 shape 与单 epoch 训练不报错。"""
"""最小冒烟测试：MoE + PINN 流程可运行。"""
from __future__ import annotations
@@ -10,82 +10,93 @@ import torch
import yaml
from src.config import load_config
from src.data import INPUT_COLUMNS
from src.data import load_raw_txt
from src.model import MLPRegressor
from src.preprocess import prepare_training_data
from src.trainer import fit
from src.model import create_model_from_config
from src.preprocess import inverse_transform_targets, prepare_training_data
from src.trainer import compute_pinn_loss, fit
def _write_synthetic_txt(path: Path, n: int = 64) -> None:
def _write_synthetic_txt(path: Path, n: int = 96) -> None:
rng = np.random.default_rng(0)
x = rng.normal(size=(n, 8))
y = np.zeros((n, 3))
y[:, 0] = rng.normal(size=n)
y[:, 1] = rng.normal(size=n)
# 第 11 列 V_pi 落在默认物理门控 [0, 500] 内
y[:, 2] = rng.uniform(1.0, 400.0, size=n)
length = np.abs(x[:, 7]) + 0.5
y[:, 0] = 2.0 - 0.2 * length + rng.normal(scale=0.05, size=n)
y[:, 1] = 0.3 + 0.1 * length + rng.normal(scale=0.03, size=n)
y[:, 2] = 20.0 / length + rng.normal(scale=0.2, size=n)
mat = np.hstack([x, y])
lines = [",".join(str(v) for v in row) for row in mat]
path.write_text("\n".join(lines), encoding="utf-8")
path.write_text("\n".join(",".join(str(v) for v in row) for row in mat), encoding="utf-8")
def test_mlp_forward_shape() -> None:
m = MLPRegressor(8, [16, 16], 3, batchnorm=False, dropout=0.0, residual=False)
def _make_cfg(tmp_path: Path, data_txt: Path, epochs: int = 1) -> Path:
cfg_dict = {
"data_path": str(data_txt),
"data": {
"test_size": 0.1,
"random_state": 123,
"filter_v_pi_max": 500.0,
},
"model": {
"input_dim": 8,
"output_dim": 3,
"hidden_dims": [16, 16],
"n_experts": 4,
"gating_hidden": 4,
"dropout_rate": 0.0,
"use_bn": False,
"activation": "relu",
},
"optimizer": {
"lr": 0.001,
"weight_decay": 0.01,
"betas": [0.9, 0.999],
},
"training": {
"batch_size": 16,
"epochs": epochs,
"num_workers": 0,
},
"physics": {
"lambda_bw_mon": 0.1,
"lambda_IL_mon": 0.1,
"lambda_vpiL": 0.05,
"lambda_smooth": 0.01,
},
"output_dir": str(tmp_path / "results"),
}
cfg_path = tmp_path / "cfg.yaml"
cfg_path.write_text(yaml.safe_dump(cfg_dict), encoding="utf-8")
return cfg_path
def test_moe_forward_shape(tmp_path: Path) -> None:
data_txt = tmp_path / "data.txt"
_write_synthetic_txt(data_txt, n=32)
cfg = load_config(_make_cfg(tmp_path, data_txt))
model = create_model_from_config(cfg)
x = torch.randn(5, 8)
y = m(x)
y = model(x)
assert y.shape == (5, 3)
def test_pinn_loss_backpropagates(tmp_path: Path) -> None:
data_txt = tmp_path / "data.txt"
_write_synthetic_txt(data_txt, n=40)
cfg = load_config(_make_cfg(tmp_path, data_txt))
model = create_model_from_config(cfg)
x = torch.randn(8, 8)
y = torch.randn(8, 3)
loss, data_loss, terms = compute_pinn_loss(model, x, y, torch.nn.MSELoss(), cfg)
loss.backward()
assert float(loss.detach()) >= float(data_loss.detach())
assert "IL_mon" in terms
assert any(p.grad is not None for p in model.parameters())
def test_one_epoch_training_pipeline(tmp_path: Path) -> None:
data_txt = tmp_path / "data.txt"
_write_synthetic_txt(data_txt, n=80)
cfg_dict = {
"data_path": str(data_txt),
"split_ratios": [0.7, 0.15, 0.15],
"random_seed": 1,
"remove_duplicate_rows": False,
"outlier_strategy": "none",
"outlier_apply_to": "targets",
"outlier_config": {
"iqr_k": 1.5,
"zscore_threshold": 4.0,
"quantile_lower": 0.001,
"quantile_upper": 0.999,
},
"remove_nonpositive_vpi": False,
"filter_v_pi_range": True,
"v_pi_min": 0.0,
"v_pi_max": 500.0,
"model": {
"input_dim": 8,
"hidden_dims": [32, 32],
"output_dim": 3,
"batchnorm": False,
"dropout": 0.0,
"residual": False,
},
"optimizer": {"name": "adamw", "lr": 0.01, "weight_decay": 0.0},
"scheduler": {
"type": "cosine",
"plateau_factor": 0.5,
"plateau_patience": 10,
"plateau_min_lr": 1e-6,
},
"training": {
"batch_size": 16,
"epochs": 1,
"early_stopping_patience": 1,
"num_workers": 0,
},
"loss": {"type": "huber", "huber_delta": 1.0, "target_weights": [1.0, 1.0, 1.0]},
"output_dir": str(tmp_path / "results"),
}
cfg_path = tmp_path / "cfg.yaml"
cfg_path.write_text(yaml.safe_dump(cfg_dict), encoding="utf-8")
cfg = load_config(cfg_path)
cfg = load_config(_make_cfg(tmp_path, data_txt, epochs=1))
df = load_raw_txt(data_txt)
run_dir = tmp_path / "run0"
@@ -95,99 +106,12 @@ def test_one_epoch_training_pipeline(tmp_path: Path) -> None:
bundle = prepare_training_data(df, cfg, run_dir)
device = torch.device("cpu")
model = MLPRegressor(
input_dim=cfg.model.input_dim,
hidden_dims=cfg.model.hidden_dims,
output_dim=cfg.model.output_dim,
batchnorm=cfg.model.batchnorm,
dropout=cfg.model.dropout,
residual=cfg.model.residual,
).to(device)
fit(model, cfg, bundle.train_loader, bundle.val_loader, run_dir, device)
assert (run_dir / "checkpoints" / "best.pt").is_file()
model = create_model_from_config(cfg).to(device)
history = fit(model, cfg, bundle, run_dir, device)
assert (run_dir / "checkpoints" / "last.pt").is_file()
assert len(history.train_loss) == 1
meta = json.loads((run_dir / "cleaning_meta.json").read_text(encoding="utf-8"))
assert meta["n_train"] > 0
def test_grouped_split_keeps_same_inputs_together(tmp_path: Path) -> None:
data_txt = tmp_path / "grouped_data.txt"
rng = np.random.default_rng(7)
rows = []
base_inputs = rng.normal(size=(24, 8))
for x in base_inputs:
for _ in range(3):
y0 = float(x[0] * 2.0 + rng.normal(scale=0.01))
y1 = float(x[1] * -1.5 + rng.normal(scale=0.01))
y2 = float(abs(x[2]) * 20.0 + 10.0 + rng.normal(scale=0.1))
rows.append(np.concatenate([x, [y0, y1, y2]]))
mat = np.asarray(rows, dtype=float)
data_txt.write_text(
"\n".join(",".join(str(v) for v in row) for row in mat),
encoding="utf-8",
)
cfg_dict = {
"data_path": str(data_txt),
"split_ratios": [0.7, 0.15, 0.15],
"random_seed": 3,
"split_mode": "grouped_stratified",
"split_stratify_target": "V_pi",
"split_stratify_bins": 6,
"remove_duplicate_rows": False,
"outlier_strategy": "none",
"outlier_apply_to": "targets",
"outlier_config": {
"iqr_k": 1.5,
"zscore_threshold": 4.0,
"quantile_lower": 0.001,
"quantile_upper": 0.999,
},
"remove_nonpositive_vpi": False,
"filter_v_pi_range": True,
"v_pi_min": 0.0,
"v_pi_max": 500.0,
"model": {
"input_dim": 8,
"hidden_dims": [16, 16],
"output_dim": 3,
"batchnorm": False,
"dropout": 0.0,
"residual": False,
},
"optimizer": {"name": "adamw", "lr": 0.01, "weight_decay": 0.0},
"scheduler": {
"type": "cosine",
"plateau_factor": 0.5,
"plateau_patience": 10,
"plateau_min_lr": 1e-6,
},
"training": {
"batch_size": 16,
"epochs": 1,
"early_stopping_patience": 1,
"num_workers": 0,
},
"loss": {"type": "huber", "huber_delta": 1.0, "target_weights": [1.0, 1.0, 1.0]},
"output_dir": str(tmp_path / "results"),
}
cfg_path = tmp_path / "cfg_grouped.yaml"
cfg_path.write_text(yaml.safe_dump(cfg_dict), encoding="utf-8")
cfg = load_config(cfg_path)
df = load_raw_txt(data_txt)
run_dir = tmp_path / "run_grouped"
run_dir.mkdir()
bundle = prepare_training_data(df, cfg, run_dir)
split_data = json.loads((run_dir / "split_indices.json").read_text(encoding="utf-8"))
split_name_by_row = {}
for split_name, indices in split_data.items():
for idx in indices:
split_name_by_row[int(idx)] = split_name
cleaned = df.reset_index(drop=True)
for _, sub in cleaned.groupby(INPUT_COLUMNS, dropna=False):
assigned = {split_name_by_row[int(i)] for i in sub.index.to_list()}
assert len(assigned) == 1
assert len(bundle.X_train) > 0
assert meta["n_test"] > 0
restored = inverse_transform_targets(bundle.y_test[:3], bundle.y_scalers)
assert restored.shape == (3, 3)