mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-07-22 05:49:55 +08:00
575 lines
20 KiB
Python
575 lines
20 KiB
Python
"""预训练脚本：在外部 LiON 数据上预训练 backbone + delivery head"""
|
||
|
||
import json
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional
|
||
|
||
import pandas as pd
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch.utils.data import DataLoader
|
||
from loguru import logger
|
||
from tqdm import tqdm
|
||
import typer
|
||
|
||
from lnp_ml.config import MODELS_DIR, PROCESSED_DATA_DIR
|
||
from lnp_ml.dataset import ExternalDeliveryDataset, collate_fn
|
||
from lnp_ml.modeling.visualization import plot_loss_curves
|
||
|
||
# MPNN ensemble 默认路径
|
||
DEFAULT_MPNN_ENSEMBLE_DIR = MODELS_DIR / "mpnn" / "all_amine_split_for_LiON"
|
||
|
||
|
||
def find_mpnn_ensemble_paths(base_dir: Path = DEFAULT_MPNN_ENSEMBLE_DIR) -> List[str]:
|
||
"""
|
||
自动查找 MPNN ensemble 的 model.pt 文件。
|
||
|
||
在 base_dir 下查找所有 cv_*/fold_*/model_*/model.pt 文件。
|
||
"""
|
||
model_paths = sorted(base_dir.glob("cv_*/fold_*/model_*/model.pt"))
|
||
if not model_paths:
|
||
raise FileNotFoundError(f"No model.pt files found in {base_dir}")
|
||
return [str(p) for p in model_paths]
|
||
from lnp_ml.modeling.models import LNPModel, LNPModelWithoutMPNN
|
||
from lnp_ml.modeling.layers.llm_prompt import DEFAULT_MOLT5_PATH
|
||
from lnp_ml.utils.seed import set_global_seed
|
||
|
||
|
||
app = typer.Typer()
|
||
|
||
|
||
class EarlyStopping:
|
||
"""早停机制"""
|
||
|
||
def __init__(self, patience: int = 10, min_delta: float = 0.0):
|
||
self.patience = patience
|
||
self.min_delta = min_delta
|
||
self.counter = 0
|
||
self.best_loss = float("inf")
|
||
|
||
def __call__(self, val_loss: float) -> bool:
|
||
if val_loss < self.best_loss - self.min_delta:
|
||
self.best_loss = val_loss
|
||
self.counter = 0
|
||
return False
|
||
self.counter += 1
|
||
return self.counter >= self.patience
|
||
|
||
|
||
def warmup_cache(model: nn.Module, smiles_list: List[str], batch_size: int = 256) -> None:
|
||
"""
|
||
预热 RDKit 特征缓存，避免训练时计算阻塞。
|
||
"""
|
||
unique_smiles = list(set(smiles_list))
|
||
logger.info(f"Warming up RDKit cache for {len(unique_smiles)} unique SMILES...")
|
||
|
||
for i in tqdm(range(0, len(unique_smiles), batch_size), desc="Cache warmup"):
|
||
batch = unique_smiles[i:i + batch_size]
|
||
model.rdkit_encoder(batch)
|
||
|
||
logger.success(f"Cache warmup complete. Cached {len(model.rdkit_encoder._cache)} SMILES.")
|
||
|
||
|
||
def train_epoch_delivery(
|
||
model: nn.Module,
|
||
loader: DataLoader,
|
||
optimizer: torch.optim.Optimizer,
|
||
device: torch.device,
|
||
epoch: int = 0,
|
||
) -> Dict[str, float]:
|
||
"""
|
||
单个 epoch 的预训练（仅 delivery 任务）。
|
||
"""
|
||
model.train()
|
||
total_loss = 0.0
|
||
n_samples = 0
|
||
|
||
pbar = tqdm(loader, desc=f"Epoch {epoch+1} [Train]", leave=False)
|
||
for batch in pbar:
|
||
smiles = batch["smiles"]
|
||
tabular = {k: v.to(device) for k, v in batch["tabular"].items()}
|
||
targets = batch["targets"]["delivery"].to(device)
|
||
mask = batch["mask"]["delivery"].to(device)
|
||
|
||
optimizer.zero_grad()
|
||
|
||
# Forward: 只预测 delivery
|
||
pred = model.forward_delivery(smiles, tabular) # [B, 1]
|
||
pred = pred.squeeze(-1) # [B]
|
||
|
||
# 计算损失（仅对有效样本）
|
||
if mask.any():
|
||
loss = nn.functional.mse_loss(pred[mask], targets[mask])
|
||
loss.backward()
|
||
optimizer.step()
|
||
total_loss += loss.item() * mask.sum().item()
|
||
n_samples += mask.sum().item()
|
||
|
||
pbar.set_postfix({"loss": total_loss / max(n_samples, 1)})
|
||
|
||
avg_loss = total_loss / max(n_samples, 1)
|
||
return {"loss": avg_loss, "n_samples": n_samples}
|
||
|
||
|
||
@torch.no_grad()
|
||
def validate_delivery(
|
||
model: nn.Module,
|
||
loader: DataLoader,
|
||
device: torch.device,
|
||
) -> Dict[str, float]:
|
||
"""
|
||
验证（仅 delivery 任务）。
|
||
"""
|
||
model.eval()
|
||
total_loss = 0.0
|
||
n_samples = 0
|
||
|
||
for batch in loader:
|
||
smiles = batch["smiles"]
|
||
tabular = {k: v.to(device) for k, v in batch["tabular"].items()}
|
||
targets = batch["targets"]["delivery"].to(device)
|
||
mask = batch["mask"]["delivery"].to(device)
|
||
|
||
pred = model.forward_delivery(smiles, tabular).squeeze(-1)
|
||
|
||
if mask.any():
|
||
loss = nn.functional.mse_loss(pred[mask], targets[mask])
|
||
total_loss += loss.item() * mask.sum().item()
|
||
n_samples += mask.sum().item()
|
||
|
||
avg_loss = total_loss / max(n_samples, 1)
|
||
return {"loss": avg_loss, "n_samples": n_samples}
|
||
|
||
|
||
def pretrain(
|
||
train_loader: DataLoader,
|
||
val_loader: DataLoader,
|
||
model: nn.Module,
|
||
device: torch.device,
|
||
lr: float = 1e-4,
|
||
weight_decay: float = 1e-5,
|
||
epochs: int = 50,
|
||
patience: int = 10,
|
||
) -> dict:
|
||
"""
|
||
预训练循环。
|
||
|
||
Returns:
|
||
训练历史和最佳验证损失
|
||
"""
|
||
model = model.to(device)
|
||
trainable_params = [p for p in model.parameters() if p.requires_grad]
|
||
optimizer = torch.optim.AdamW(trainable_params, lr=lr, weight_decay=weight_decay)
|
||
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
||
optimizer, mode="min", factor=0.5, patience=5
|
||
)
|
||
early_stopping = EarlyStopping(patience=patience)
|
||
|
||
history = {"train": [], "val": []}
|
||
best_val_loss = float("inf")
|
||
best_state = None
|
||
|
||
for epoch in range(epochs):
|
||
# Train
|
||
train_metrics = train_epoch_delivery(model, train_loader, optimizer, device, epoch)
|
||
|
||
# Validate
|
||
val_metrics = validate_delivery(model, val_loader, device)
|
||
|
||
# Log
|
||
logger.info(
|
||
f"Epoch {epoch + 1}/{epochs} | "
|
||
f"Train Loss: {train_metrics['loss']:.4f} | "
|
||
f"Val Loss: {val_metrics['loss']:.4f}"
|
||
)
|
||
|
||
history["train"].append(train_metrics)
|
||
history["val"].append(val_metrics)
|
||
|
||
# Learning rate scheduling
|
||
scheduler.step(val_metrics["loss"])
|
||
|
||
# Save best model
|
||
if val_metrics["loss"] < best_val_loss:
|
||
best_val_loss = val_metrics["loss"]
|
||
best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
|
||
logger.info(f" -> New best model (val_loss={best_val_loss:.4f})")
|
||
|
||
# Early stopping
|
||
if early_stopping(val_metrics["loss"]):
|
||
logger.info(f"Early stopping at epoch {epoch + 1}")
|
||
break
|
||
|
||
# Restore best model
|
||
if best_state is not None:
|
||
model.load_state_dict(best_state)
|
||
|
||
return {
|
||
"history": history,
|
||
"best_val_loss": best_val_loss,
|
||
}
|
||
|
||
|
||
@app.command()
|
||
def main(
|
||
# 数据路径（已处理的 parquet 文件）
|
||
train_path: Path = PROCESSED_DATA_DIR / "train_pretrain.parquet",
|
||
val_path: Path = PROCESSED_DATA_DIR / "val_pretrain.parquet",
|
||
output_dir: Path = MODELS_DIR,
|
||
# 模型参数
|
||
d_model: int = 256,
|
||
num_heads: int = 8,
|
||
n_attn_layers: int = 4,
|
||
set_transformer_block: str = "sab", # "sab" | "isab"
|
||
fusion_strategy: str = "attention",
|
||
head_hidden_dim: int = 128,
|
||
dropout: float = 0.1,
|
||
# MoE 参数（消融开关）
|
||
use_moe: bool = False,
|
||
moe_n_experts: int = 4,
|
||
moe_top_k: int = 2,
|
||
moe_expert_hidden_mult: int = 2,
|
||
moe_jitter_noise: float = 0.0,
|
||
# LLM 参数（消融开关）
|
||
use_llm: bool = False,
|
||
llm_model_path: str = DEFAULT_MOLT5_PATH,
|
||
llm_freeze: bool = True,
|
||
llm_use_lora: bool = False,
|
||
llm_lora_r: int = 8,
|
||
llm_lora_alpha: int = 16,
|
||
llm_lora_dropout: float = 0.05,
|
||
# MPNN 参数（可选）
|
||
use_mpnn: bool = False, # 启用 MPNN，自动从默认路径加载 ensemble
|
||
mpnn_checkpoint: Optional[str] = None,
|
||
mpnn_ensemble_paths: Optional[str] = None,
|
||
mpnn_device: str = "cpu",
|
||
# 训练参数
|
||
batch_size: int = 64,
|
||
lr: float = 1e-4,
|
||
weight_decay: float = 1e-5,
|
||
epochs: int = 50,
|
||
patience: int = 10,
|
||
# 随机种子
|
||
seed: int = 42,
|
||
# 设备
|
||
device: str = "cuda" if torch.cuda.is_available() else "cpu",
|
||
):
|
||
"""
|
||
在外部 LiON 数据上预训练 backbone + delivery head。
|
||
|
||
需要先运行 `make data_pretrain` 生成 parquet 文件。
|
||
|
||
使用 --use-mpnn 启用 MPNN encoder（自动从 models/mpnn/all_amine_split_for_LiON 加载）。
|
||
|
||
产出:
|
||
- models/pretrain_delivery.pt: 包含 backbone + delivery head 权重
|
||
- models/pretrain_history.json: 训练历史
|
||
"""
|
||
set_global_seed(seed)
|
||
logger.info(f"Using device: {device} | seed: {seed}")
|
||
device_obj = torch.device(device)
|
||
|
||
# 加载已处理的 parquet 文件
|
||
logger.info(f"Loading train data from {train_path}")
|
||
train_df = pd.read_parquet(train_path)
|
||
train_dataset = ExternalDeliveryDataset(train_df)
|
||
|
||
logger.info(f"Loading val data from {val_path}")
|
||
val_df = pd.read_parquet(val_path)
|
||
val_dataset = ExternalDeliveryDataset(val_df)
|
||
|
||
logger.info(f"Train samples: {len(train_dataset)}, Val samples: {len(val_dataset)}")
|
||
|
||
train_loader = DataLoader(
|
||
train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn
|
||
)
|
||
val_loader = DataLoader(
|
||
val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
|
||
)
|
||
|
||
# 解析 MPNN 配置
|
||
# 优先级：mpnn_checkpoint > mpnn_ensemble_paths > use_mpnn（自动查找）
|
||
ensemble_paths_list = None
|
||
if mpnn_ensemble_paths:
|
||
ensemble_paths_list = mpnn_ensemble_paths.split(",")
|
||
elif use_mpnn and mpnn_checkpoint is None:
|
||
# --use-mpnn 但没有指定具体路径，自动查找
|
||
logger.info(f"Auto-detecting MPNN ensemble from {DEFAULT_MPNN_ENSEMBLE_DIR}")
|
||
ensemble_paths_list = find_mpnn_ensemble_paths()
|
||
logger.info(f"Found {len(ensemble_paths_list)} MPNN models")
|
||
|
||
enable_mpnn = mpnn_checkpoint is not None or ensemble_paths_list is not None
|
||
|
||
# 创建模型
|
||
logger.info(
|
||
f"Creating model (use_mpnn={enable_mpnn}, use_moe={use_moe}, use_llm={use_llm})..."
|
||
)
|
||
common_kwargs = dict(
|
||
d_model=d_model,
|
||
num_heads=num_heads,
|
||
n_attn_layers=n_attn_layers,
|
||
set_transformer_block=set_transformer_block,
|
||
fusion_strategy=fusion_strategy,
|
||
head_hidden_dim=head_hidden_dim,
|
||
dropout=dropout,
|
||
use_moe=use_moe,
|
||
moe_n_experts=moe_n_experts,
|
||
moe_top_k=moe_top_k,
|
||
moe_expert_hidden_mult=moe_expert_hidden_mult,
|
||
moe_jitter_noise=moe_jitter_noise,
|
||
use_llm=use_llm,
|
||
llm_model_path=llm_model_path,
|
||
llm_freeze=llm_freeze,
|
||
llm_use_lora=llm_use_lora,
|
||
llm_lora_r=llm_lora_r,
|
||
llm_lora_alpha=llm_lora_alpha,
|
||
llm_lora_dropout=llm_lora_dropout,
|
||
)
|
||
if enable_mpnn:
|
||
model = LNPModel(
|
||
mpnn_checkpoint=mpnn_checkpoint,
|
||
mpnn_ensemble_paths=ensemble_paths_list,
|
||
mpnn_device=mpnn_device,
|
||
**common_kwargs,
|
||
)
|
||
else:
|
||
model = LNPModelWithoutMPNN(**common_kwargs)
|
||
|
||
n_params_total = sum(p.numel() for p in model.parameters())
|
||
n_params_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
logger.info(f"Model parameters: {n_params_total:,} total, {n_params_trainable:,} trainable")
|
||
|
||
# 预热 RDKit 缓存（避免训练时阻塞）
|
||
all_smiles = train_df["smiles"].tolist() + val_df["smiles"].tolist()
|
||
warmup_cache(model, all_smiles, batch_size=256)
|
||
|
||
# 预训练
|
||
logger.info("Starting pretraining on external data (delivery only)...")
|
||
result = pretrain(
|
||
train_loader=train_loader,
|
||
val_loader=val_loader,
|
||
model=model,
|
||
device=device_obj,
|
||
lr=lr,
|
||
weight_decay=weight_decay,
|
||
epochs=epochs,
|
||
patience=patience,
|
||
)
|
||
|
||
# 保存预训练 checkpoint
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
checkpoint_path = output_dir / "pretrain_delivery.pt"
|
||
torch.save(
|
||
{
|
||
"model_state_dict": model.state_dict(),
|
||
"backbone_state_dict": model.get_backbone_state_dict(),
|
||
"delivery_head_state_dict": model.get_delivery_head_state_dict(),
|
||
"config": {
|
||
"d_model": d_model,
|
||
"num_heads": num_heads,
|
||
"n_attn_layers": n_attn_layers,
|
||
"set_transformer_block": set_transformer_block,
|
||
"fusion_strategy": fusion_strategy,
|
||
"head_hidden_dim": head_hidden_dim,
|
||
"dropout": dropout,
|
||
"use_mpnn": enable_mpnn,
|
||
"use_moe": use_moe,
|
||
"moe_n_experts": moe_n_experts,
|
||
"moe_top_k": moe_top_k,
|
||
"moe_expert_hidden_mult": moe_expert_hidden_mult,
|
||
"moe_jitter_noise": moe_jitter_noise,
|
||
"use_llm": use_llm,
|
||
"llm_model_path": llm_model_path,
|
||
"llm_freeze": llm_freeze,
|
||
"llm_use_lora": llm_use_lora,
|
||
"llm_lora_r": llm_lora_r,
|
||
"llm_lora_alpha": llm_lora_alpha,
|
||
"llm_lora_dropout": llm_lora_dropout,
|
||
},
|
||
"best_val_loss": result["best_val_loss"],
|
||
},
|
||
checkpoint_path,
|
||
)
|
||
logger.success(f"Saved pretrain checkpoint to {checkpoint_path}")
|
||
|
||
# 保存训练历史
|
||
history_path = output_dir / "pretrain_history.json"
|
||
with open(history_path, "w") as f:
|
||
json.dump(result["history"], f, indent=2)
|
||
logger.success(f"Saved pretrain history to {history_path}")
|
||
|
||
# 绘制 loss 曲线图
|
||
loss_plot_path = output_dir / "pretrain_loss_curves.png"
|
||
plot_loss_curves(
|
||
history=result["history"],
|
||
output_path=loss_plot_path,
|
||
title="Pretrain Loss Curves (Delivery)",
|
||
)
|
||
logger.success(f"Saved loss curves plot to {loss_plot_path}")
|
||
|
||
logger.success(
|
||
f"Pretraining complete! Best val_loss: {result['best_val_loss']:.4f}"
|
||
)
|
||
|
||
|
||
@app.command()
|
||
def test(
|
||
# 数据路径
|
||
val_path: Path = PROCESSED_DATA_DIR / "val_pretrain.parquet",
|
||
model_path: Path = MODELS_DIR / "pretrain_delivery.pt",
|
||
output_path: Path = MODELS_DIR / "pretrain_test_results.json",
|
||
# MPNN 参数
|
||
use_mpnn: bool = False,
|
||
mpnn_device: str = "cpu",
|
||
# 其他参数
|
||
batch_size: int = 64,
|
||
device: str = "cuda" if torch.cuda.is_available() else "cpu",
|
||
):
|
||
"""
|
||
评估 pretrain 模型在外部数据上的 delivery 预测性能。
|
||
|
||
输出详细指标：MSE, RMSE, MAE, R²
|
||
"""
|
||
import numpy as np
|
||
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
|
||
|
||
logger.info(f"Using device: {device}")
|
||
device_obj = torch.device(device)
|
||
|
||
# 加载模型
|
||
logger.info(f"Loading pretrain model from {model_path}")
|
||
checkpoint = torch.load(model_path, map_location=device_obj, weights_only=False)
|
||
config = checkpoint["config"]
|
||
|
||
# 解析 MPNN 配置
|
||
enable_mpnn = config.get("use_mpnn", False)
|
||
common_kwargs = dict(
|
||
d_model=config["d_model"],
|
||
num_heads=config["num_heads"],
|
||
n_attn_layers=config["n_attn_layers"],
|
||
set_transformer_block=config.get("set_transformer_block", "sab"),
|
||
fusion_strategy=config["fusion_strategy"],
|
||
head_hidden_dim=config["head_hidden_dim"],
|
||
dropout=config["dropout"],
|
||
use_moe=config.get("use_moe", False),
|
||
moe_n_experts=config.get("moe_n_experts", 4),
|
||
moe_top_k=config.get("moe_top_k", 2),
|
||
moe_expert_hidden_mult=config.get("moe_expert_hidden_mult", 2),
|
||
moe_jitter_noise=config.get("moe_jitter_noise", 0.0),
|
||
use_llm=config.get("use_llm", False),
|
||
llm_model_path=config.get("llm_model_path", DEFAULT_MOLT5_PATH),
|
||
llm_freeze=config.get("llm_freeze", True),
|
||
llm_use_lora=config.get("llm_use_lora", False),
|
||
llm_lora_r=config.get("llm_lora_r", 8),
|
||
llm_lora_alpha=config.get("llm_lora_alpha", 16),
|
||
llm_lora_dropout=config.get("llm_lora_dropout", 0.05),
|
||
)
|
||
if enable_mpnn or use_mpnn:
|
||
logger.info(f"Auto-detecting MPNN ensemble from {DEFAULT_MPNN_ENSEMBLE_DIR}")
|
||
ensemble_paths = find_mpnn_ensemble_paths()
|
||
logger.info(f"Found {len(ensemble_paths)} MPNN models")
|
||
model = LNPModel(
|
||
mpnn_ensemble_paths=ensemble_paths,
|
||
mpnn_device=mpnn_device,
|
||
**common_kwargs,
|
||
)
|
||
else:
|
||
model = LNPModelWithoutMPNN(**common_kwargs)
|
||
|
||
model.load_state_dict(checkpoint["model_state_dict"], strict=False)
|
||
model.to(device_obj)
|
||
model.eval()
|
||
|
||
logger.info(f"Model config: {config}")
|
||
logger.info(f"Best val_loss from training: {checkpoint.get('best_val_loss', 'N/A')}")
|
||
|
||
# 加载数据
|
||
logger.info(f"Loading validation data from {val_path}")
|
||
val_df = pd.read_parquet(val_path)
|
||
val_dataset = ExternalDeliveryDataset(val_df)
|
||
val_loader = DataLoader(
|
||
val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
|
||
)
|
||
logger.info(f"Validation samples: {len(val_dataset)}")
|
||
|
||
# 预测
|
||
logger.info("Running predictions...")
|
||
all_preds = []
|
||
all_targets = []
|
||
all_masks = []
|
||
|
||
with torch.no_grad():
|
||
for batch in tqdm(val_loader, desc="Predicting"):
|
||
smiles = batch["smiles"]
|
||
tabular = {k: v.to(device_obj) for k, v in batch["tabular"].items()}
|
||
targets = batch["targets"]["delivery"].numpy()
|
||
mask = batch["mask"]["delivery"].numpy()
|
||
|
||
pred = model.forward_delivery(smiles, tabular).squeeze(-1).cpu().numpy()
|
||
|
||
all_preds.extend(pred)
|
||
all_targets.extend(targets)
|
||
all_masks.extend(mask)
|
||
|
||
# 转为数组
|
||
all_preds = np.array(all_preds)
|
||
all_targets = np.array(all_targets)
|
||
all_masks = np.array(all_masks, dtype=bool)
|
||
|
||
# 只计算有效样本
|
||
y_pred = all_preds[all_masks]
|
||
y_true = all_targets[all_masks]
|
||
|
||
# 计算指标
|
||
mse = float(mean_squared_error(y_true, y_pred))
|
||
rmse = float(np.sqrt(mse))
|
||
mae = float(mean_absolute_error(y_true, y_pred))
|
||
r2 = float(r2_score(y_true, y_pred))
|
||
|
||
# 额外统计
|
||
correlation = float(np.corrcoef(y_true, y_pred)[0, 1])
|
||
|
||
results = {
|
||
"model_path": str(model_path),
|
||
"val_path": str(val_path),
|
||
"n_samples": int(all_masks.sum()),
|
||
"metrics": {
|
||
"mse": mse,
|
||
"rmse": rmse,
|
||
"mae": mae,
|
||
"r2": r2,
|
||
"correlation": correlation,
|
||
},
|
||
"statistics": {
|
||
"y_true_mean": float(y_true.mean()),
|
||
"y_true_std": float(y_true.std()),
|
||
"y_pred_mean": float(y_pred.mean()),
|
||
"y_pred_std": float(y_pred.std()),
|
||
}
|
||
}
|
||
|
||
# 打印结果
|
||
logger.info("\n" + "=" * 50)
|
||
logger.info("PRETRAIN MODEL EVALUATION (Delivery)")
|
||
logger.info("=" * 50)
|
||
logger.info(f"Samples: {results['n_samples']}")
|
||
logger.info("\n[Metrics]")
|
||
logger.info(f" MSE: {mse:.4f}")
|
||
logger.info(f" RMSE: {rmse:.4f}")
|
||
logger.info(f" MAE: {mae:.4f}")
|
||
logger.info(f" R²: {r2:.4f}")
|
||
logger.info(f" Correlation: {correlation:.4f}")
|
||
logger.info("\n[Statistics]")
|
||
logger.info(f" True: mean={y_true.mean():.4f}, std={y_true.std():.4f}")
|
||
logger.info(f" Pred: mean={y_pred.mean():.4f}, std={y_pred.std():.4f}")
|
||
|
||
# 保存结果
|
||
with open(output_path, "w") as f:
|
||
json.dump(results, f, indent=2)
|
||
logger.success(f"\nSaved results to {output_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
app()
|
||
|