#!/usr/bin/env python3 """Train one XGBoost PKS assistant model without reading the test set.""" from __future__ import annotations import argparse import importlib import json from collections import defaultdict from pathlib import Path import joblib import numpy as np from sklearn.impute import SimpleImputer from sklearn.metrics import average_precision_score, precision_recall_curve from sklearn.neighbors import KNeighborsClassifier from sklearn.pipeline import Pipeline from sklearn.preprocessing import StandardScaler from xgboost import XGBClassifier from xgboost.callback import TrainingCallback def feature_module(model_type: str): """Pressure and vibration use independent feature modules.""" return importlib.import_module( "generalized_features" if model_type == "pressure" else "generalized_features_vibration" ) class ProgressCallback(TrainingCallback): """Print round, learning rate, and train/validation loss for each boosting round.""" def __init__(self, rounds: int, learning_rate: float) -> None: self.rounds = rounds self.learning_rate = learning_rate def after_iteration(self, model, epoch: int, evals_log): losses = [] for dataset, metrics in evals_log.items(): for name, values in metrics.items(): if values: losses.append(f"{dataset}-{name}={values[-1]:.5f}") print( f"[round {epoch + 1:>4}/{self.rounds}] lr={self.learning_rate:.4f} " + " ".join(losses), flush=True, ) return False def load_samples(directory: Path) -> list[dict]: return [json.loads(path.read_text(encoding="utf-8")) for path in sorted(directory.glob("*.json"))] def matrix(feature_dicts: list[dict], names: list[str]) -> np.ndarray: return np.asarray([[item.get(name) for name in names] for item in feature_dicts], dtype=float) def stratified_normal_sample(samples: list[dict], max_ratio: int, seed: int) -> list[dict]: positives = [sample for sample in samples if sample["label"] == 1] normals = [sample for sample in samples if sample["label"] == 0] limit = len(positives) * max_ratio if len(normals) <= limit: return samples buckets = defaultdict(list) for sample in normals: month = sample["start_hour"][:7] buckets[(sample["unit"], sample["batch"], month)].append(sample) rng = np.random.default_rng(seed) selected = [] # Round-robin quotas preserve unit/month coverage rather than favoring the # longest uninterrupted running cycle. keys = sorted(buckets) quota, remainder = divmod(limit, len(keys)) for index, key in enumerate(keys): bucket = buckets[key] take = min(len(bucket), quota + (1 if index < remainder else 0)) selected.extend(bucket[index] for index in rng.choice(len(bucket), size=take, replace=False)) if len(selected) < limit: leftovers = [sample for sample in normals if sample not in selected] need = min(limit - len(selected), len(leftovers)) selected.extend(leftovers[index] for index in rng.choice(len(leftovers), size=need, replace=False)) return positives + selected def choose_threshold(y: np.ndarray, probabilities: np.ndarray) -> tuple[float, dict]: precision, recall, thresholds = precision_recall_curve(y, probabilities) f1 = 2 * precision * recall / np.maximum(precision + recall, 1e-12) index = int(np.nanargmax(f1[:-1])) if len(thresholds) else 0 threshold = float(thresholds[index]) if len(thresholds) else 0.5 return threshold, { "threshold_selection": "max_validation_f1", "threshold": threshold, "validation_precision": float(precision[index]), "validation_recall": float(recall[index]), "validation_f1": float(f1[index]), "validation_pr_auc": float(average_precision_score(y, probabilities)), } def main() -> None: parser = argparse.ArgumentParser(description="训练 PKS 压力或振动模型，不读取测试集") parser.add_argument("model_type", choices=("pressure", "vibration")) parser.add_argument("--datasets-dir", default=str(Path(__file__).resolve().parent / "datasets")) parser.add_argument("--models-dir", default=str(Path(__file__).resolve().parent / "models")) sampling = parser.add_mutually_exclusive_group() sampling.add_argument("--normal-ratio", type=int, default=None) sampling.add_argument("--all-normal", action="store_true", help="使用全部正常训练样本，不下采样、不复制") parser.add_argument("--seed", type=int, default=20260921) parser.add_argument("--rounds", type=int, default=250) parser.add_argument("--early-stopping-rounds", type=int, default=30) args = parser.parse_args() if args.normal_ratio is not None and args.normal_ratio < 1: parser.error("--normal-ratio 必须大于等于 1") if args.rounds < 1 or args.early_stopping_rounds < 1: parser.error("--rounds 和 --early-stopping-rounds 必须大于等于 1") ratio = args.normal_ratio if args.normal_ratio is not None else (4 if args.model_type == "pressure" else 5) gf = feature_module(args.model_type) dataset = Path(args.datasets_dir) / args.model_type train_all = load_samples(dataset / "train") validation = load_samples(dataset / "validation") if not train_all or not validation: parser.error("训练或验证样本为空") if not any(sample["label"] == 1 for sample in train_all) or not any(sample["label"] == 1 for sample in validation): parser.error("训练和验证集都必须包含异常样本") train = train_all if args.all_normal else stratified_normal_sample(train_all, ratio, args.seed) baselines = gf.compute_unit_baselines(train_all) names = gf.feature_names() x_train, y_train = matrix(gf.feature_matrix(train, baselines), names), np.asarray([sample["label"] for sample in train]) x_validation, y_validation = matrix(gf.feature_matrix(validation, baselines), names), np.asarray([sample["label"] for sample in validation]) imputer = SimpleImputer(strategy="median") x_train = imputer.fit_transform(x_train) x_validation = imputer.transform(x_validation) scale_pos_weight = float(np.sum(y_train == 0) / max(np.sum(y_train == 1), 1)) learning_rate = 0.05 print( f"training: rounds={args.rounds}, learning_rate={learning_rate:.4f}, " f"early_stopping_rounds={args.early_stopping_rounds}, samples={len(train):,}", flush=True, ) model = XGBClassifier( n_estimators=args.rounds, max_depth=3, learning_rate=learning_rate, subsample=0.8, colsample_bytree=0.8, min_child_weight=3, objective="binary:logistic", eval_metric="logloss", scale_pos_weight=scale_pos_weight, random_state=args.seed, n_jobs=4, early_stopping_rounds=args.early_stopping_rounds, callbacks=[ProgressCallback(args.rounds, learning_rate)], ) eval_set = [(x_train, y_train), (x_validation, y_validation)] model.fit(x_train, y_train, eval_set=eval_set, verbose=False) rounds_run = len(model.evals_result()["validation_0"]["logloss"]) stopped_early = rounds_run < args.rounds print( f"training complete: rounds_run={rounds_run}/{args.rounds}, " f"best_iteration={model.best_iteration + 1}, best_validation_logloss={float(model.best_score):.5f}, " f"early_stopped={'yes' if stopped_early else 'no'}", flush=True, ) validation_probabilities = model.predict_proba(x_validation)[:, 1] threshold, validation_metrics = choose_threshold(y_validation, validation_probabilities) similarity = Pipeline([("scale", StandardScaler()), ("knn", KNeighborsClassifier(n_neighbors=min(5, len(train)), weights="distance"))]) similarity.fit(x_train, y_train) models = Path(args.models_dir) models.mkdir(parents=True, exist_ok=True) # The progress callback is defined in this script (module __main__); drop it # before serialization so the artifact can be loaded by other entry points. model.callbacks = [] joblib.dump( {"model": model, "imputer": imputer, "knn": similarity, "feature_names": names, "unit_baselines": baselines}, models / f"{args.model_type}_model.joblib", ) (models / f"{args.model_type}_unit_baselines.json").write_text( json.dumps(baselines, ensure_ascii=False, indent=2), encoding="utf-8" ) metadata = { "model_type": args.model_type, "algorithm": "XGBoost binary classifier + KNN similarity", "feature_version": "generalized_cycle_relative_v2", "feature_names": names, "dropped_absolute_features": list(gf.ABSOLUTE_FEATURES), "unit_baselines": baselines, "normal_sampling_strategy": "all_normal" if args.all_normal else "stratified_max_ratio", "normal_to_anomaly_training_ratio_limit": None if args.all_normal else ratio, "actual_normal_to_anomaly_training_ratio": float(np.sum(y_train == 0) / max(np.sum(y_train == 1), 1)), "training_samples_before_sampling": len(train_all), "training_samples_after_sampling": len(train), "training_normal": int(np.sum(y_train == 0)), "training_anomaly": int(np.sum(y_train == 1)), "scale_pos_weight": scale_pos_weight, "validation_samples": len(validation), "test_set_read": False, **validation_metrics, "rounds_requested": args.rounds, "early_stopping_rounds_requested": args.early_stopping_rounds, "rounds_run": rounds_run, "early_stopped": stopped_early, "best_iteration": int(model.best_iteration), "best_score": float(model.best_score), } (models / f"{args.model_type}_metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(metadata, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()