model_training_lightgbm.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. import lightgbm as lgb
  2. import numpy as np
  3. from sklearn.model_selection import train_test_split
  4. from sklearn.metrics import mean_squared_error,mean_absolute_error
  5. from flask import Flask,request
  6. import time
  7. import traceback
  8. import logging
  9. from common.database_dml import get_data_from_mongo,insert_pickle_model_into_mongo
  10. from common.processing_data_common import missing_features,str_to_list
  11. app = Flask('model_training_lightgbm——service')
  12. def build_model(df,args):
  13. np.random.seed(42)
  14. #lightgbm预测下
  15. numerical_features,categorical_features,label,model_name,num_boost_round,model_params,col_time = str_to_list(args['numerical_features']),str_to_list(args['categorical_features']),args['label'],args['model_name'],int(args['num_boost_round']),eval(args['model_params']),args['col_time']
  16. features = numerical_features+categorical_features
  17. print("features:************",features)
  18. if 'is_limit' in df.columns:
  19. df = df[df['is_limit']==False]
  20. # 清洗特征平均缺失率大于20%的天
  21. df = missing_features(df, features, col_time)
  22. df = df[~np.isnan(df[label])]
  23. # 拆分数据为训练集和测试集
  24. X_train, X_test, y_train, y_test = train_test_split(df[features], df[label], test_size=0.2, random_state=42)
  25. # 创建LightGBM数据集
  26. lgb_train = lgb.Dataset(X_train, y_train,categorical_feature=categorical_features)
  27. lgb_eval = lgb.Dataset(X_test, y_test, reference=lgb_train)
  28. # 设置参数
  29. params = {
  30. 'objective': 'regression',
  31. 'metric': 'rmse',
  32. 'boosting_type': 'gbdt',
  33. 'verbose':1
  34. }
  35. params.update(model_params)
  36. # 训练模型
  37. print('Starting training...')
  38. gbm = lgb.train(params,
  39. lgb_train,
  40. num_boost_round=num_boost_round,
  41. valid_sets=[lgb_train, lgb_eval],
  42. )
  43. y_pred = gbm.predict(X_test, num_iteration=gbm.best_iteration)
  44. # 评估
  45. mse = mean_squared_error(y_test, y_pred)
  46. rmse = np.sqrt(mse)
  47. mae = mean_absolute_error(y_test, y_pred)
  48. print(f'The test rmse is: {rmse},"The test mae is:"{mae}')
  49. return gbm
  50. @app.route('/model_training_lightgbm', methods=['POST'])
  51. def model_training_lightgbm():
  52. # 获取程序开始时间
  53. start_time = time.time()
  54. result = {}
  55. success = 0
  56. print("Program starts execution!")
  57. try:
  58. args = request.values.to_dict()
  59. print('args',args)
  60. logger.info(args)
  61. power_df = get_data_from_mongo(args)
  62. model = build_model(power_df,args)
  63. insert_pickle_model_into_mongo(model,args)
  64. success = 1
  65. except Exception as e:
  66. my_exception = traceback.format_exc()
  67. my_exception.replace("\n","\t")
  68. result['msg'] = my_exception
  69. end_time = time.time()
  70. result['success'] = success
  71. result['args'] = args
  72. result['start_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))
  73. result['end_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))
  74. print("Program execution ends!")
  75. return result
  76. if __name__=="__main__":
  77. print("Program starts execution!")
  78. logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
  79. logger = logging.getLogger("model_training_lightgbm log")
  80. from waitress import serve
  81. serve(app, host="0.0.0.0", port=10089)
  82. print("server start!")