algorithm_platform/models_processing/model_tf/tf_bp_train.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :tf_bp_train.py
# @Time      :2025/2/13 13:35
# @Author    :David
# @Company: shenyang JY

import json, copy
import numpy as np
from flask import Flask, request
import traceback
import logging, argparse
from data_processing.data_operation.data_handler import DataHandler
import time, yaml
from models_processing.model_tf.tf_bp import BPHandler
from common.database_dml_koi import *
import matplotlib.pyplot as plt
from common.logs import Log
logger = Log('tf_bp').logger
np.random.seed(42)  # NumPy随机种子
app = Flask('tf_bp_train——service')

with app.app_context():
    current_dir = os.path.dirname(os.path.abspath(__file__))
    with open(os.path.join(current_dir, 'bp.yaml'), 'r', encoding='utf-8') as f:
        args = yaml.safe_load(f)
    dh = DataHandler(logger, args)
    bp = BPHandler(logger, args)

@app.before_request
def update_config():
    # ------------ 整理参数，整合请求参数 ------------
    args_dict = request.values.to_dict()
    args_dict['features'] = args_dict['features'].split(',')
    args.update(args_dict)
    opt = argparse.Namespace(**args)
    dh.opt = opt
    bp.opt = opt
    logger.info(args)

@app.route('/nn_bp_training', methods=['POST'])
def model_training_bp():
    # 获取程序开始时间
    start_time = time.time()
    result = {}
    success = 0
    print("Program starts execution!")
    try:
        # ------------ 获取数据，预处理训练数据 ------------
        train_data = get_data_from_mongo(args)
        train_x, train_y, valid_x, valid_y, scaled_train_bytes, scaled_target_bytes, scaled_cap = dh.train_data_handler(train_data, bp_data=True)
        bp.opt.cap = round(scaled_cap, 2)
        bp.opt.Model['input_size'] = len(dh.opt.features)
        # ------------ 训练模型 ------------
        # 1. 如果是加强训练模式，先加载预训练模型特征参数，再预处理训练数据
        # 2. 如果是普通模式，先预处理训练数据，再根据训练数据特征加载模型
        model = bp.train_init() if bp.opt.Model['add_train'] else bp.get_keras_model(bp.opt)
        if bp.opt.Model['add_train']:
            if model:
                feas = json.loads(bp.model_params).get('features', args['features'])
                if set(feas).issubset(set(dh.opt.features)):
                    dh.opt.features = list(feas)
                    train_x, train_y, valid_x, valid_y, scaled_train_bytes, scaled_target_bytes, scaled_cap = dh.train_data_handler(train_data)
                else:
                    model = bp.get_keras_model(bp.opt)
                    logger.info("训练数据特征，不满足，加强训练模型特征")
            else:
                model = bp.get_keras_model(bp.opt)
        bp_model = bp.training(model, [train_x, train_y, valid_x, valid_y])
        # ------------ 保存模型 ------------
        args['Model']['features'] = ','.join(dh.opt.features)
        args['params'] = json.dumps(args)
        args['descr'] = '测试'
        args['gen_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
        insert_trained_model_into_mongo(bp_model, args)
        insert_scaler_model_into_mongo(scaled_train_bytes, scaled_target_bytes, args)
        success = 1
    except Exception as e:
        my_exception = traceback.format_exc()
        my_exception.replace("\n", "\t")
        result['msg'] = my_exception
    end_time = time.time()
    result['success'] = success
    result['args'] = args
    result['start_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))
    result['end_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))
    print("Program execution ends!")
    return result


if __name__ == "__main__":
    print("Program starts execution!")
    from waitress import serve
    serve(
        app,
        host="0.0.0.0",
        port=10111,
        threads=8,  # 指定线程数（默认4，根据硬件调整）
        channel_timeout=600  # 连接超时时间（秒）
    )
    # print("server start!")
    # args_dict = {"mongodb_database": 'david_test', 'scaler_table': 'j00083_scaler', 'model_name': 'bp1.0.test',
    # 'model_table': 'j00083_model', 'mongodb_read_table': 'j00083', 'col_time': 'dateTime',
    # 'features': 'speed10,direction10,speed30,direction30,speed50,direction50,speed70,direction70,speed90,direction90,speed110,direction110,speed150,direction150,speed170,direction170'}
    # args_dict['features'] = args_dict['features'].split(',')
    # arguments.update(args_dict)
    # dh = DataHandler(logger, arguments)
    # bp = BPHandler(logger)
    # opt = argparse.Namespace(**arguments)
    # opt.Model['input_size'] = len(opt.features)
    # train_data = get_data_from_mongo(args_dict)
    # train_x, valid_x, train_y, valid_y, scaled_train_bytes, scaled_target_bytes = dh.train_data_handler(train_data, opt, bp_data=True)
    # bp_model = bp.training(opt, [train_x, train_y, valid_x, valid_y])
    #
    # args_dict['gen_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
    # args_dict['params'] = arguments
    # args_dict['descr'] = '测试'
    # insert_trained_model_into_mongo(bp_model, args_dict)
    # insert_scaler_model_into_mongo(scaled_train_bytes, scaled_target_bytes, args_dict)