CoolFace
Apppublic

Qionk/a-share-quant

sourceHugging Faceupdated 4mo agoView on Hugging Face
0likes
quick_train_worker.py33 linesDownload Raw Back to root
1"""快速训练 worker — 独立进程,不受 Streamlit 干扰"""2import os, sys3 4os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"5os.environ["CUDA_VISIBLE_DEVICES"] = ""6 7# 必须在 import akshare/pandas 之前先 import TF,否则线程池冲突导致死锁8import tensorflow  # noqa: F4019 10import pickle11 12if __name__ == "__main__":13    input_path = sys.argv[1]14    output_path = sys.argv[2]15 16    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))17 18    with open(input_path, "rb") as f:19        args = pickle.load(f)20 21    from src.predict.training import train_all_models22 23    results = train_all_models(24        args["df"], args["selected_models"], args["config"],25        forecast_days=args["forecast_days"], quick=True)26 27    for r in results.values():28        r.model_object = None29        r.scaler = None30 31    with open(output_path, "wb") as f:32        pickle.dump(results, f)33