#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 负荷预测模型训练脚本 功能：启动训练任务并监控进度 """ import requests import argparse import time import json from datetime import datetime # API基础地址 API_BASE_URL = "http://47.96.16.84:5005" def start_training(learning_rate=0.001, batch_size=32, max_epochs=100, window_size=96, hidden_size=128, num_layers=2, split_ratio='811'): """ 启动训练任务 """ url = f"{API_BASE_URL}/api/load-forecast/train/start" data = { "task_name": f"负荷预测训练_{datetime.now().strftime('%Y%m%d_%H%M%S')}", "model_type": "cnn_rbilstm", "config": { "learning_rate": learning_rate, "batch_size": batch_size, "epochs": max_epochs, "window_size": window_size, "hidden_size": hidden_size, "num_layers": num_layers, "split_ratio": split_ratio } } try: response = requests.post(url, json=data) response.raise_for_status() result = response.json() if result.get('code') == 200: task_id = result['data']['task_id'] print(f"训练任务启动成功，任务ID: {task_id}") return task_id else: print(f"启动训练失败: {result.get('error')}") return None except Exception as e: print(f"启动训练异常: {e}") return None def get_training_status(task_id): """ 获取训练状态 """ url = f"{API_BASE_URL}/api/load-forecast/train/status/{task_id}" try: response = requests.get(url) response.raise_for_status() result = response.json() if result.get('code') == 200: return result['data'] else: print(f"获取训练状态失败: {result.get('error')}") return None except Exception as e: print(f"获取训练状态异常: {e}") return None def get_training_logs(task_id): """ 获取训练日志 """ url = f"{API_BASE_URL}/api/load-forecast/train/logs/{task_id}" try: response = requests.get(url) response.raise_for_status() result = response.json() if result.get('code') == 200: return result['data']['logs'] else: print(f"获取训练日志失败: {result.get('error')}") return [] except Exception as e: print(f"获取训练日志异常: {e}") return [] def monitor_training(task_id, check_interval=10, max_wait_time=3600): """ 监控训练进度 """ start_time = time.time() last_log_count = 0 while time.time() - start_time < max_wait_time: status = get_training_status(task_id) if not status: time.sleep(check_interval) continue # 打印状态信息 print(f"\r轮次: {status['progress']['current_epoch']}/{status['progress']['total_epochs']} | " f"进度: {status['progress']['percentage']:.1f}% | " f"损失: {status['metrics']['train_loss']:.4f} | " f"验证损失: {status['metrics']['val_loss']:.4f} | " f"耗时: {status['elapsed_time']}s", end="") # 打印最新日志 logs = get_training_logs(task_id) if len(logs) > last_log_count: new_logs = logs[last_log_count:] for log in new_logs: print(f"\n{log}") last_log_count = len(logs) # 检查训练是否完成 if status['status'] in ['completed', 'failed', 'stopped']: print(f"\n训练状态: {status['status']}") if 'model_path' in status and status['model_path']: print(f"模型保存路径: {status['model_path']}") return status['status'] time.sleep(check_interval) print("\n训练监控超时") return 'timeout' def train_and_monitor(learning_rate=0.001, batch_size=32, max_epochs=100, window_size=96, hidden_size=128, num_layers=2, split_ratio='811', check_interval=10, max_wait_time=3600): """ 启动训练并监控进度 """ task_id = start_training(learning_rate, batch_size, max_epochs, window_size, hidden_size, num_layers, split_ratio) if not task_id: return print(f"开始监控训练进度 (任务ID: {task_id})") monitor_training(task_id, check_interval, max_wait_time) def check_status(task_id): """ 检查训练状态 """ status = get_training_status(task_id) if status: print("训练状态:") print(f"任务ID: {task_id}") print(f"状态: {status['status']}") print(f"进度: {status['progress']['current_epoch']}/{status['progress']['total_epochs']} ({status['progress']['percentage']:.1f}%)") print(f"损失: {status['metrics']['train_loss']:.4f}") print(f"验证损失: {status['metrics']['val_loss']:.4f}") print(f"耗时: {status['elapsed_time']}秒") if 'model_path' in status and status['model_path']: print(f"模型路径: {status['model_path']}") else: print("无法获取训练状态") def main(): parser = argparse.ArgumentParser(description='负荷预测模型训练脚本') parser.add_argument('--action', type=str, required=True, choices=['train', 'train_and_monitor', 'check_status'], help='操作类型') parser.add_argument('--task_id', type=str, help='任务ID') parser.add_argument('--learning_rate', type=float, default=0.001, help='学习率') parser.add_argument('--batch_size', type=int, default=32, help='批大小') parser.add_argument('--max_epochs', type=int, default=100, help='最大训练轮次') parser.add_argument('--window_size', type=int, default=96, help='窗口大小') parser.add_argument('--hidden_size', type=int, default=128, help='隐藏层大小') parser.add_argument('--num_layers', type=int, default=2, help='层数') parser.add_argument('--split_ratio', type=str, default='811', help='数据分割比例') parser.add_argument('--check_interval', type=int, default=10, help='状态检查间隔（秒）') parser.add_argument('--max_wait_time', type=int, default=3600, help='最大等待时间（秒）') args = parser.parse_args() if args.action == 'train': start_training( learning_rate=args.learning_rate, batch_size=args.batch_size, max_epochs=args.max_epochs, window_size=args.window_size, hidden_size=args.hidden_size, num_layers=args.num_layers, split_ratio=args.split_ratio ) elif args.action == 'train_and_monitor': train_and_monitor( learning_rate=args.learning_rate, batch_size=args.batch_size, max_epochs=args.max_epochs, window_size=args.window_size, hidden_size=args.hidden_size, num_layers=args.num_layers, split_ratio=args.split_ratio, check_interval=args.check_interval, max_wait_time=args.max_wait_time ) elif args.action == 'check_status': if not args.task_id: print('错误：必须指定任务ID') return check_status(args.task_id) if __name__ == '__main__': main()