257 lines
8.8 KiB
Python
257 lines
8.8 KiB
Python
|
"""
|
|||
|
|
本地训练脚本 — 生成合成数据、训练 MLP、导出 ONNX
|
|||
|
|
运行方式（在 backend/ 目录下）：
|
|||
|
|
python train_models.py
|
|||
|
|
|
|||
|
|
依赖（仅本地训练用，不进 Docker）：
|
|||
|
|
pip install torch onnx onnxruntime scikit-learn numpy
|
|||
|
|
"""
|
|||
|
|
import sys, json, time
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.optim as optim
|
|||
|
|
from torch.utils.data import DataLoader, TensorDataset
|
|||
|
|
|
|||
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|||
|
|
from app.services.prediction import AcidSpeedModel, TensionModel, QualityPredictionModel
|
|||
|
|
|
|||
|
|
PT_DIR = Path(__file__).parent / "app" / "services" / "pt_models"
|
|||
|
|
PT_DIR.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
SEED = 2024
|
|||
|
|
N = 12000
|
|||
|
|
np.random.seed(SEED)
|
|||
|
|
torch.manual_seed(SEED)
|
|||
|
|
|
|||
|
|
TENSION_ZONES = [
|
|||
|
|
"inlet", "s1_roller", "acid_entry",
|
|||
|
|
"acid1", "acid2", "acid3",
|
|||
|
|
"rinse", "leveler", "s2_roller", "outlet",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─── 网络结构 ───────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
class MLP(nn.Module):
|
|||
|
|
def __init__(self, in_dim: int, out_dim: int, hidden=(128, 64, 32)):
|
|||
|
|
super().__init__()
|
|||
|
|
layers: list = []
|
|||
|
|
prev = in_dim
|
|||
|
|
for h in hidden:
|
|||
|
|
layers += [nn.Linear(prev, h), nn.ReLU()]
|
|||
|
|
prev = h
|
|||
|
|
layers.append(nn.Linear(prev, out_dim))
|
|||
|
|
self.net = nn.Sequential(*layers)
|
|||
|
|
|
|||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|||
|
|
return self.net(x)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─── 训练通用函数 ───────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def fit(model: nn.Module, X: np.ndarray, y: np.ndarray,
|
|||
|
|
epochs=300, lr=1e-3, batch_size=512) -> nn.Module:
|
|||
|
|
Xt = torch.from_numpy(X)
|
|||
|
|
yt = torch.from_numpy(y)
|
|||
|
|
dl = DataLoader(TensorDataset(Xt, yt), batch_size=batch_size, shuffle=True)
|
|||
|
|
opt = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
|
|||
|
|
sched = optim.lr_scheduler.CosineAnnealingLR(opt, epochs)
|
|||
|
|
loss_fn = nn.MSELoss()
|
|||
|
|
|
|||
|
|
model.train()
|
|||
|
|
for ep in range(1, epochs + 1):
|
|||
|
|
tot = 0.0
|
|||
|
|
for xb, yb in dl:
|
|||
|
|
opt.zero_grad()
|
|||
|
|
loss = loss_fn(model(xb), yb)
|
|||
|
|
loss.backward()
|
|||
|
|
opt.step()
|
|||
|
|
tot += loss.item() * len(xb)
|
|||
|
|
sched.step()
|
|||
|
|
if ep % 100 == 0:
|
|||
|
|
print(f" ep {ep:3d}/{epochs} RMSE={((tot/len(Xt))**0.5):.5f}")
|
|||
|
|
return model
|
|||
|
|
|
|||
|
|
|
|||
|
|
def z_scale(arr: np.ndarray, mean=None, std=None):
|
|||
|
|
if mean is None:
|
|||
|
|
mean = arr.mean(axis=0)
|
|||
|
|
std = arr.std(axis=0) + 1e-8
|
|||
|
|
return ((arr - mean) / std).astype(np.float32), mean, std
|
|||
|
|
|
|||
|
|
|
|||
|
|
def export_onnx(model: nn.Module, in_dim: int, path: Path):
|
|||
|
|
model.eval()
|
|||
|
|
dummy = torch.zeros(1, in_dim)
|
|||
|
|
torch.onnx.export(
|
|||
|
|
model, dummy, str(path),
|
|||
|
|
input_names=["input"], output_names=["output"],
|
|||
|
|
dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
|
|||
|
|
opset_version=17,
|
|||
|
|
)
|
|||
|
|
print(f" → {path.name} ({path.stat().st_size//1024} KB)")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─── 1. 酸洗速度模型 ────────────────────────────────────────────────────────
|
|||
|
|
# 输入(14): thickness, scale_weight, conc×6, temp×6
|
|||
|
|
# 输出(1): max_speed
|
|||
|
|
|
|||
|
|
def gen_acid_speed(n: int):
|
|||
|
|
rng = np.random.default_rng(SEED)
|
|||
|
|
Xs, ys = [], []
|
|||
|
|
skip = 0
|
|||
|
|
while len(Xs) < n:
|
|||
|
|
t = rng.uniform(0.5, 8.0)
|
|||
|
|
sw = rng.uniform(4.0, 18.0)
|
|||
|
|
conc = rng.uniform(60, 240, 6).tolist()
|
|||
|
|
temp = rng.uniform(52, 87, 6).tolist()
|
|||
|
|
tpi = rng.uniform(88, 97)
|
|||
|
|
try:
|
|||
|
|
m = AcidSpeedModel(
|
|||
|
|
thickness=t, width=1000.0, steel_grade="Q235",
|
|||
|
|
acid_conc_list=conc, acid_temp_list=temp,
|
|||
|
|
scale_weight=sw, target_pi=tpi,
|
|||
|
|
)
|
|||
|
|
spd = float(m.calculate()["max_speed"])
|
|||
|
|
except Exception:
|
|||
|
|
skip += 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 模拟真实工况偏差：±6% 相对噪声 + 钢种系数扰动
|
|||
|
|
steel_factor = rng.choice([0.92, 0.96, 1.00, 1.03, 1.06])
|
|||
|
|
noise = rng.normal(1.0, 0.06)
|
|||
|
|
spd_n = float(np.clip(spd * noise * steel_factor, 20, 180))
|
|||
|
|
|
|||
|
|
Xs.append([t, sw] + conc + temp)
|
|||
|
|
ys.append([spd_n])
|
|||
|
|
|
|||
|
|
print(f" acid_speed: {len(Xs)} samples (skipped {skip})")
|
|||
|
|
return np.array(Xs, np.float32), np.array(ys, np.float32)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─── 2. 张力模型 ────────────────────────────────────────────────────────────
|
|||
|
|
# 输入(4): thickness, width, yield_strength, tension_coef
|
|||
|
|
# 输出(10): 10 区段张力 kN
|
|||
|
|
|
|||
|
|
def gen_tension(n: int):
|
|||
|
|
rng = np.random.default_rng(SEED + 1)
|
|||
|
|
Xs, ys = [], []
|
|||
|
|
while len(Xs) < n:
|
|||
|
|
t = rng.uniform(0.5, 8.0)
|
|||
|
|
w = rng.uniform(600, 1600)
|
|||
|
|
ys_ = rng.uniform(150, 600)
|
|||
|
|
tc = rng.uniform(0.15, 0.35)
|
|||
|
|
|
|||
|
|
m = TensionModel(thickness=t, width=w, yield_strength=ys_, tension_coef=tc)
|
|||
|
|
res = m.calculate()
|
|||
|
|
tensions = [res["zones"][z]["tension_kN"] for z in TENSION_ZONES]
|
|||
|
|
|
|||
|
|
# 各区段独立噪声（实测张力传感器精度约 ±4%）
|
|||
|
|
noise = rng.normal(1.0, 0.04, 10)
|
|||
|
|
tensions_n = [float(np.clip(v * noise[i], 0.1, 9999)) for i, v in enumerate(tensions)]
|
|||
|
|
|
|||
|
|
Xs.append([t, w, ys_, tc])
|
|||
|
|
ys.append(tensions_n)
|
|||
|
|
|
|||
|
|
print(f" tension: {len(Xs)} samples")
|
|||
|
|
return np.array(Xs, np.float32), np.array(ys, np.float32)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─── 3. 质量预测模型 ─────────────────────────────────────────────────────────
|
|||
|
|
# 输入(6): thickness, avg_speed, acid_conc_avg, acid_temp_avg, scale_weight, fe_conc_avg
|
|||
|
|
# 输出(2): pi_score, surface_score
|
|||
|
|
|
|||
|
|
def gen_quality(n: int):
|
|||
|
|
rng = np.random.default_rng(SEED + 2)
|
|||
|
|
Xs, ys = [], []
|
|||
|
|
while len(Xs) < n:
|
|||
|
|
t = rng.uniform(0.5, 8.0)
|
|||
|
|
spd = rng.uniform(20, 180)
|
|||
|
|
conc = rng.uniform(60, 240)
|
|||
|
|
temp = rng.uniform(50, 90)
|
|||
|
|
sw = rng.uniform(4, 18)
|
|||
|
|
fe = rng.uniform(20, 130)
|
|||
|
|
|
|||
|
|
m = QualityPredictionModel(
|
|||
|
|
thickness=t, avg_speed=spd,
|
|||
|
|
acid_conc_avg=conc, acid_temp_avg=temp,
|
|||
|
|
scale_weight=sw, fe_conc_avg=fe,
|
|||
|
|
)
|
|||
|
|
res = m.calculate()
|
|||
|
|
pi = res["pi_score"]
|
|||
|
|
suf = res["surface_score"]
|
|||
|
|
|
|||
|
|
# ±6% 噪声模拟质检测量不确定度
|
|||
|
|
pi_n = float(np.clip(pi * rng.normal(1.0, 0.06), 0, 100))
|
|||
|
|
suf_n = float(np.clip(suf * rng.normal(1.0, 0.06), 0, 100))
|
|||
|
|
|
|||
|
|
Xs.append([t, spd, conc, temp, sw, fe])
|
|||
|
|
ys.append([pi_n, suf_n])
|
|||
|
|
|
|||
|
|
print(f" quality: {len(Xs)} samples")
|
|||
|
|
return np.array(Xs, np.float32), np.array(ys, np.float32)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─── 主流程 ─────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
scalers: dict = {}
|
|||
|
|
t0 = time.time()
|
|||
|
|
|
|||
|
|
# ── 酸洗速度 ──
|
|||
|
|
print("\n[1/3] 酸洗速度模型")
|
|||
|
|
X, y = gen_acid_speed(N)
|
|||
|
|
Xn, Xm, Xs = z_scale(X)
|
|||
|
|
yn, ym, ys_ = z_scale(y)
|
|||
|
|
model = MLP(14, 1, hidden=(128, 64, 32))
|
|||
|
|
print(" 训练中...")
|
|||
|
|
fit(model, Xn, yn, epochs=300)
|
|||
|
|
export_onnx(model, 14, PT_DIR / "acid_speed.onnx")
|
|||
|
|
scalers["acid_speed"] = {
|
|||
|
|
"X_mean": Xm.tolist(), "X_std": Xs.tolist(),
|
|||
|
|
"y_mean": ym.tolist(), "y_std": ys_.tolist(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# ── 张力 ──
|
|||
|
|
print("\n[2/3] 张力模型")
|
|||
|
|
X, y = gen_tension(N)
|
|||
|
|
Xn, Xm, Xs = z_scale(X)
|
|||
|
|
yn, ym, ys_ = z_scale(y)
|
|||
|
|
model = MLP(4, 10, hidden=(64, 64, 32))
|
|||
|
|
print(" 训练中...")
|
|||
|
|
fit(model, Xn, yn, epochs=300)
|
|||
|
|
export_onnx(model, 4, PT_DIR / "tension.onnx")
|
|||
|
|
scalers["tension"] = {
|
|||
|
|
"X_mean": Xm.tolist(), "X_std": Xs.tolist(),
|
|||
|
|
"y_mean": ym.tolist(), "y_std": ys_.tolist(),
|
|||
|
|
"zone_names": TENSION_ZONES,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# ── 质量 ──
|
|||
|
|
print("\n[3/3] 质量预测模型")
|
|||
|
|
X, y = gen_quality(N)
|
|||
|
|
Xn, Xm, Xs = z_scale(X)
|
|||
|
|
yn, ym, ys_ = z_scale(y)
|
|||
|
|
model = MLP(6, 2, hidden=(64, 32))
|
|||
|
|
print(" 训练中...")
|
|||
|
|
fit(model, Xn, yn, epochs=300)
|
|||
|
|
export_onnx(model, 6, PT_DIR / "quality.onnx")
|
|||
|
|
scalers["quality"] = {
|
|||
|
|
"X_mean": Xm.tolist(), "X_std": Xs.tolist(),
|
|||
|
|
"y_mean": ym.tolist(), "y_std": ys_.tolist(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# ── 保存 scaler 参数 ──
|
|||
|
|
scaler_path = PT_DIR / "scalers.json"
|
|||
|
|
with open(scaler_path, "w") as f:
|
|||
|
|
json.dump(scalers, f, indent=2)
|
|||
|
|
print(f"\n scalers → {scaler_path.name}")
|
|||
|
|
print(f"\n完成 ({time.time()-t0:.1f}s)\n")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|