634 lines
24 KiB
Python
634 lines
24 KiB
Python
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()
|
||
|