Files
timelib/run_gui.py

634 lines
24 KiB
Python
Raw Normal View History

2026-01-29 17:08:54 +08:00
import argparse
import os
import sys
import torch
import torch.backends
from exp.exp_long_term_forecasting import Exp_Long_Term_Forecast
from exp.exp_imputation import Exp_Imputation
from exp.exp_short_term_forecasting import Exp_Short_Term_Forecast
from exp.exp_anomaly_detection import Exp_Anomaly_Detection
from exp.exp_classification import Exp_Classification
from exp.exp_zero_shot_forecasting import Exp_Zero_Shot_Forecast
from utils.print_args import print_args
import random
import numpy as np
import threading
import queue
import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
import matplotlib
matplotlib.use('TkAgg')
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import re
import io
from contextlib import redirect_stdout
class QueueWriter:
def __init__(self, queue):
self.queue = queue
def write(self, text):
self.queue.put(text)
def flush(self):
pass
class TrainingThread(threading.Thread):
def __init__(self, args, log_queue, loss_queue, dataset_info_queue):
super().__init__()
self.args = args
self.log_queue = log_queue
self.loss_queue = loss_queue
self.dataset_info_queue = dataset_info_queue
self.daemon = True
def run(self):
# 重定向stdout到队列
original_stdout = sys.stdout
queue_writer = QueueWriter(self.log_queue)
sys.stdout = queue_writer
try:
fix_seed = 2021
random.seed(fix_seed)
torch.manual_seed(fix_seed)
np.random.seed(fix_seed)
# 设备设置
if torch.cuda.is_available() and self.args.use_gpu:
self.args.device = torch.device('cuda:{}'.format(self.args.gpu))
else:
if hasattr(torch.backends, "mps"):
self.args.device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
else:
self.args.device = torch.device("cpu")
if self.args.use_gpu and self.args.use_multi_gpu:
self.args.devices = self.args.devices.replace(' ', '')
device_ids = self.args.devices.split(',')
self.args.device_ids = [int(id_) for id_ in device_ids]
self.args.gpu = self.args.device_ids[0]
# 选择实验类
if self.args.task_name == 'long_term_forecast':
Exp = Exp_Long_Term_Forecast
elif self.args.task_name == 'short_term_forecast':
Exp = Exp_Short_Term_Forecast
elif self.args.task_name == 'imputation':
Exp = Exp_Imputation
elif self.args.task_name == 'anomaly_detection':
Exp = Exp_Anomaly_Detection
elif self.args.task_name == 'classification':
Exp = Exp_Classification
elif self.args.task_name == 'zero_shot_forecast':
Exp = Exp_Zero_Shot_Forecast
else:
Exp = Exp_Long_Term_Forecast
# 创建实验实例
exp = Exp(self.args)
# 获取数据集信息
try:
train_data, train_loader = exp._get_data(flag='train')
test_data, test_loader = exp._get_data(flag='test')
val_data, val_loader = exp._get_data(flag='val')
train_size = len(train_data)
test_size = len(test_data)
val_size = len(val_data) if val_data else 0
dataset_info = {
'train_size': train_size,
'test_size': test_size,
'val_size': val_size,
'train_batches': len(train_loader),
'test_batches': len(test_loader),
'val_batches': len(val_loader) if val_loader else 0
}
self.dataset_info_queue.put(dataset_info)
except Exception as e:
self.log_queue.put(f"获取数据集信息失败: {str(e)}\n")
if self.args.is_training:
for ii in range(self.args.itr):
setting = '{}_{}_{}_{}_ft{}_sl{}_ll{}_pl{}_dm{}_nh{}_el{}_dl{}_df{}_expand{}_dc{}_fc{}_eb{}_dt{}_{}_{}'.format(
self.args.task_name,
self.args.model_id,
self.args.model,
self.args.data,
self.args.features,
self.args.seq_len,
self.args.label_len,
self.args.pred_len,
self.args.d_model,
self.args.n_heads,
self.args.e_layers,
self.args.d_layers,
self.args.d_ff,
self.args.expand,
self.args.d_conv,
self.args.factor,
self.args.embed,
self.args.distil,
self.args.des, ii)
print(f'开始训练: {setting}')
# 训练过程
exp.train(setting)
print(f'开始测试: {setting}')
exp.test(setting)
if self.args.gpu_type == 'mps':
torch.backends.mps.empty_cache()
elif self.args.gpu_type == 'cuda':
torch.cuda.empty_cache()
print("训练完成!")
except Exception as e:
self.log_queue.put(f"训练过程中出错: {str(e)}\n")
import traceback
self.log_queue.put(traceback.format_exc() + "\n")
finally:
sys.stdout = original_stdout
class TrainingGUI:
def __init__(self, root):
self.root = root
self.root.title("时间序列预测模型训练界面")
self.root.geometry("1400x900")
# 训练线程
self.training_thread = None
self.is_training = False
# 队列
self.log_queue = queue.Queue()
self.loss_queue = queue.Queue()
self.dataset_info_queue = queue.Queue()
# Loss数据
self.train_losses = []
self.val_losses = []
self.test_losses = []
self.epochs = []
self.create_widgets()
self.process_queue()
def create_widgets(self):
# 创建主框架
main_frame = ttk.Frame(self.root, padding="10")
main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
self.root.columnconfigure(0, weight=1)
self.root.rowconfigure(0, weight=1)
main_frame.columnconfigure(1, weight=1)
main_frame.rowconfigure(1, weight=1)
# 左侧参数面板
left_frame = ttk.Frame(main_frame)
left_frame.grid(row=0, column=0, rowspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), padx=(0, 10))
left_frame.columnconfigure(1, weight=1)
# 创建滚动框架
canvas = tk.Canvas(left_frame, width=400, height=700)
scrollbar = ttk.Scrollbar(left_frame, orient="vertical", command=canvas.yview)
scrollable_frame = ttk.Frame(canvas)
scrollable_frame.bind(
"<Configure>",
lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
)
canvas.create_window((0, 0), window=scrollable_frame, anchor="nw")
canvas.configure(yscrollcommand=scrollbar.set)
canvas.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
scrollbar.grid(row=0, column=1, sticky=(tk.N, tk.S))
left_frame.rowconfigure(0, weight=1)
# 参数输入
self.params = {}
row = 0
# 基本配置
ttk.Label(scrollable_frame, text="基本配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
row += 1
self.create_param_row(scrollable_frame, "任务类型 (task_name)", "long_term_forecast", row)
row += 1
self.create_param_row(scrollable_frame, "模型 (model)", "TimesNet", row)
row += 1
self.create_param_row(scrollable_frame, "模型ID (model_id)", "test", row)
row += 1
# 数据配置
ttk.Label(scrollable_frame, text="数据配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
row += 1
self.create_param_row(scrollable_frame, "数据集 (data)", "ETTh1", row)
row += 1
self.create_param_row(scrollable_frame, "数据路径 (root_path)", "./dataset/ETT-small/", row)
row += 1
self.create_param_row(scrollable_frame, "数据文件 (data_path)", "ETTh1.csv", row)
row += 1
self.create_param_row(scrollable_frame, "特征类型 (features)", "M", row)
row += 1
self.create_param_row(scrollable_frame, "频率 (freq)", "h", row)
row += 1
# 预测配置
ttk.Label(scrollable_frame, text="预测配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
row += 1
self.create_param_row(scrollable_frame, "序列长度 (seq_len)", "96", row, int)
row += 1
self.create_param_row(scrollable_frame, "标签长度 (label_len)", "48", row, int)
row += 1
self.create_param_row(scrollable_frame, "预测长度 (pred_len)", "96", row, int)
row += 1
# 模型配置
ttk.Label(scrollable_frame, text="模型配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
row += 1
self.create_param_row(scrollable_frame, "模型维度 (d_model)", "512", row, int)
row += 1
self.create_param_row(scrollable_frame, "注意力头数 (n_heads)", "8", row, int)
row += 1
self.create_param_row(scrollable_frame, "编码器层数 (e_layers)", "2", row, int)
row += 1
self.create_param_row(scrollable_frame, "解码器层数 (d_layers)", "1", row, int)
row += 1
self.create_param_row(scrollable_frame, "FFN维度 (d_ff)", "2048", row, int)
row += 1
self.create_param_row(scrollable_frame, "Dropout", "0.1", row, float)
row += 1
# 训练配置
ttk.Label(scrollable_frame, text="训练配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
row += 1
self.create_param_row(scrollable_frame, "训练轮数 (train_epochs)", "10", row, int)
row += 1
self.create_param_row(scrollable_frame, "批次大小 (batch_size)", "32", row, int)
row += 1
self.create_param_row(scrollable_frame, "学习率 (learning_rate)", "0.0001", row, float)
row += 1
self.create_param_row(scrollable_frame, "早停耐心值 (patience)", "3", row, int)
row += 1
self.create_param_row(scrollable_frame, "实验次数 (itr)", "1", row, int)
row += 1
# GPU配置
ttk.Label(scrollable_frame, text="GPU配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
row += 1
self.create_param_row(scrollable_frame, "使用GPU (use_gpu)", "True", row, bool)
row += 1
self.create_param_row(scrollable_frame, "GPU设备 (gpu)", "0", row, int)
row += 1
# 按钮
button_frame = ttk.Frame(scrollable_frame)
button_frame.grid(row=row, column=0, columnspan=2, pady=20)
row += 1
self.train_button = ttk.Button(button_frame, text="开始训练", command=self.start_training, width=20)
self.train_button.pack(side=tk.LEFT, padx=5)
self.stop_button = ttk.Button(button_frame, text="停止训练", command=self.stop_training, width=20, state=tk.DISABLED)
self.stop_button.pack(side=tk.LEFT, padx=5)
# 右侧显示面板
right_frame = ttk.Frame(main_frame)
right_frame.grid(row=0, column=1, sticky=(tk.W, tk.E, tk.N, tk.S))
right_frame.columnconfigure(0, weight=1)
right_frame.rowconfigure(1, weight=1)
# 数据集信息
info_frame = ttk.LabelFrame(right_frame, text="数据集信息", padding="10")
info_frame.grid(row=0, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
info_frame.columnconfigure(0, weight=1)
self.dataset_info_text = tk.Text(info_frame, height=6, wrap=tk.WORD)
self.dataset_info_text.pack(fill=tk.BOTH, expand=True)
self.dataset_info_text.insert("1.0", "等待训练开始...\n")
self.dataset_info_text.config(state=tk.DISABLED)
# 创建标签页
notebook = ttk.Notebook(right_frame)
notebook.grid(row=1, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
# 日志标签页
log_frame = ttk.Frame(notebook)
notebook.add(log_frame, text="训练日志")
log_frame.columnconfigure(0, weight=1)
log_frame.rowconfigure(0, weight=1)
self.log_text = scrolledtext.ScrolledText(log_frame, wrap=tk.WORD, height=20)
self.log_text.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
# Loss可视化标签页
loss_frame = ttk.Frame(notebook)
notebook.add(loss_frame, text="Loss可视化")
loss_frame.columnconfigure(0, weight=1)
loss_frame.rowconfigure(0, weight=1)
self.fig = Figure(figsize=(10, 6), dpi=100)
self.ax = self.fig.add_subplot(111)
self.ax.set_xlabel('Epoch')
self.ax.set_ylabel('Loss')
self.ax.set_title('训练过程Loss曲线')
self.ax.grid(True)
self.canvas = FigureCanvasTkAgg(self.fig, loss_frame)
self.canvas.get_tk_widget().grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
def create_param_row(self, parent, label, default, row, param_type=str):
ttk.Label(parent, text=label, width=20).grid(row=row, column=0, sticky=tk.W, padx=5, pady=2)
entry = ttk.Entry(parent, width=25)
entry.insert(0, str(default))
entry.grid(row=row, column=1, sticky=(tk.W, tk.E), padx=5, pady=2)
# 提取参数名:如果标签包含括号,提取括号内的内容;否则使用标签的小写形式
if '(' in label and ')' in label:
param_name = label.split('(')[1].split(')')[0]
else:
# 没有括号的标签,使用标签的小写形式作为参数名
param_name = label.lower()
self.params[param_name] = (entry, param_type)
def get_args(self):
args = argparse.Namespace()
# 基本配置
args.task_name = self.params['task_name'][0].get()
args.model = self.params['model'][0].get()
args.model_id = self.params['model_id'][0].get()
args.is_training = 1
# 数据配置
args.data = self.params['data'][0].get()
args.root_path = self.params['root_path'][0].get()
args.data_path = self.params['data_path'][0].get()
args.features = self.params['features'][0].get()
args.target = 'OT'
args.freq = self.params['freq'][0].get()
args.checkpoints = './checkpoints/'
# 预测配置
args.seq_len = int(self.params['seq_len'][0].get())
args.label_len = int(self.params['label_len'][0].get())
args.pred_len = int(self.params['pred_len'][0].get())
args.seasonal_patterns = 'Monthly'
args.inverse = False
# 输入任务
args.mask_rate = 0.25
# 异常检测任务
args.anomaly_ratio = 0.25
# 模型配置
args.expand = 2
args.d_conv = 4
args.top_k = 5
args.num_kernels = 6
args.enc_in = 7
args.dec_in = 7
args.c_out = 7
args.d_model = int(self.params['d_model'][0].get())
args.n_heads = int(self.params['n_heads'][0].get())
args.e_layers = int(self.params['e_layers'][0].get())
args.d_layers = int(self.params['d_layers'][0].get())
args.d_ff = int(self.params['d_ff'][0].get())
args.moving_avg = 25
args.factor = 1
args.distil = True
args.dropout = float(self.params['dropout'][0].get())
args.embed = 'timeF'
args.activation = 'gelu'
args.channel_independence = 1
args.decomp_method = 'moving_avg'
args.use_norm = 1
args.down_sampling_layers = 0
args.down_sampling_window = 1
args.down_sampling_method = None
args.seg_len = 96
# 优化配置
args.num_workers = 10
args.itr = int(self.params['itr'][0].get())
args.train_epochs = int(self.params['train_epochs'][0].get())
args.batch_size = int(self.params['batch_size'][0].get())
args.patience = int(self.params['patience'][0].get())
args.learning_rate = float(self.params['learning_rate'][0].get())
args.des = 'test'
args.loss = 'MSE'
args.lradj = 'type1'
args.use_amp = False
# GPU配置
use_gpu_str = self.params['use_gpu'][0].get().lower()
args.use_gpu = use_gpu_str in ['true', '1', 'yes']
args.gpu = int(self.params['gpu'][0].get())
args.gpu_type = 'cuda'
args.use_multi_gpu = False
args.devices = '0,1,2,3'
# 其他参数
args.p_hidden_dims = [128, 128]
args.p_hidden_layers = 2
args.use_dtw = False
args.augmentation_ratio = 0
args.seed = 2
args.jitter = False
args.scaling = False
args.permutation = False
args.randompermutation = False
args.magwarp = False
args.timewarp = False
args.windowslice = False
args.windowwarp = False
args.rotation = False
args.spawner = False
args.dtwwarp = False
args.shapedtwwarp = False
args.wdba = False
args.discdtw = False
args.discsdtw = False
args.extra_tag = ""
args.patch_len = 16
args.node_dim = 10
args.gcn_depth = 2
args.gcn_dropout = 0.3
args.propalpha = 0.3
args.conv_channel = 32
args.skip_channel = 32
args.individual = False
args.alpha = 0.1
args.top_p = 0.5
args.pos = 1
return args
def start_training(self):
if self.is_training:
messagebox.showwarning("警告", "训练已在进行中!")
return
try:
args = self.get_args()
self.is_training = True
self.train_button.config(state=tk.DISABLED)
self.stop_button.config(state=tk.NORMAL)
# 清空日志和loss
self.log_text.delete("1.0", tk.END)
self.train_losses = []
self.val_losses = []
self.test_losses = []
self.epochs = []
self.ax.clear()
self.ax.set_xlabel('Epoch')
self.ax.set_ylabel('Loss')
self.ax.set_title('训练过程Loss曲线')
self.ax.grid(True)
self.canvas.draw()
# 启动训练线程
self.training_thread = TrainingThread(args, self.log_queue, self.loss_queue, self.dataset_info_queue)
self.training_thread.start()
self.log_text.insert(tk.END, "训练线程已启动...\n")
self.log_text.see(tk.END)
except Exception as e:
messagebox.showerror("错误", f"启动训练失败: {str(e)}")
self.is_training = False
self.train_button.config(state=tk.NORMAL)
self.stop_button.config(state=tk.DISABLED)
def stop_training(self):
# 注意:这个实现可能无法立即停止训练
# 更好的实现需要修改训练代码以支持中断
if self.training_thread and self.training_thread.is_alive():
messagebox.showwarning("警告", "训练无法立即停止请等待当前epoch完成")
# 可以设置一个标志位,让训练代码检查
def process_queue(self):
# 处理日志队列
try:
while True:
msg = self.log_queue.get_nowait()
self.log_text.insert(tk.END, msg)
self.log_text.see(tk.END)
# 检测训练完成
if "训练完成" in msg or "训练过程中出错" in msg:
if self.training_thread and not self.training_thread.is_alive():
self.is_training = False
self.train_button.config(state=tk.NORMAL)
self.stop_button.config(state=tk.DISABLED)
# 从日志中提取loss信息
# 匹配格式: "Epoch: {0}, Steps: {1} | Train Loss: {2:.7f} Vali Loss: {3:.7f} Test Loss: {4:.7f}"
loss_pattern = r'Epoch:\s+(\d+).*?Train Loss:\s+([\d.]+).*?Vali Loss:\s+([\d.]+).*?Test Loss:\s+([\d.]+)'
match = re.search(loss_pattern, msg)
if match:
epoch = int(match.group(1))
train_loss = float(match.group(2))
val_loss = float(match.group(3))
test_loss = float(match.group(4))
self.loss_queue.put({
'epoch': epoch,
'train_loss': train_loss,
'val_loss': val_loss,
'test_loss': test_loss
})
except queue.Empty:
pass
# 检查训练线程是否结束
if self.is_training and self.training_thread and not self.training_thread.is_alive():
self.is_training = False
self.train_button.config(state=tk.NORMAL)
self.stop_button.config(state=tk.DISABLED)
# 处理数据集信息队列
try:
while True:
info = self.dataset_info_queue.get_nowait()
self.dataset_info_text.config(state=tk.NORMAL)
self.dataset_info_text.delete("1.0", tk.END)
self.dataset_info_text.insert("1.0",
f"训练集大小: {info['train_size']}\n"
f"测试集大小: {info['test_size']}\n"
f"验证集大小: {info['val_size']}\n"
f"训练批次数: {info['train_batches']}\n"
f"测试批次数: {info['test_batches']}\n"
f"验证批次数: {info['val_batches']}\n"
)
self.dataset_info_text.config(state=tk.DISABLED)
except queue.Empty:
pass
# 处理loss队列
try:
while True:
loss_data = self.loss_queue.get_nowait()
epoch = loss_data.get('epoch', 0)
train_loss = loss_data.get('train_loss', 0)
val_loss = loss_data.get('val_loss', 0)
test_loss = loss_data.get('test_loss', 0)
self.epochs.append(epoch)
self.train_losses.append(train_loss)
self.val_losses.append(val_loss)
self.test_losses.append(test_loss)
self.update_loss_plot()
except queue.Empty:
pass
self.root.after(100, self.process_queue)
def update_loss_plot(self):
if len(self.epochs) > 0:
self.ax.clear()
self.ax.set_xlabel('Epoch')
self.ax.set_ylabel('Loss')
self.ax.set_title('训练过程Loss曲线')
self.ax.grid(True)
if len(self.train_losses) > 0:
self.ax.plot(self.epochs, self.train_losses, label='Train Loss', marker='o')
if len(self.val_losses) > 0:
self.ax.plot(self.epochs, self.val_losses, label='Val Loss', marker='s')
if len(self.test_losses) > 0:
self.ax.plot(self.epochs, self.test_losses, label='Test Loss', marker='^')
self.ax.legend()
self.canvas.draw()
def main():
root = tk.Tk()
app = TrainingGUI(root)
root.mainloop()
if __name__ == '__main__':
main()