lib/dl_other_dataset.py

@@ -1,31 +0,0 @@
-if __name__ == '__main__':
-    import os
-    import subprocess
-
-    # Kaggle 数据集列表
-    datasets = {
-        "hangzhou_taxi": "changyuheng/hz-taxi",
-        "nyc_taxi": "new-york-city/nyc-taxi-trip-duration",
-        "hangzhou_bike": "changyuheng/hz-bike"
-    }
-
-    # 下载保存目录
-    save_dir = "./datasets"
-    os.makedirs(save_dir, exist_ok=True)
-
-    # 检查 Kaggle API 配置
-    kaggle_json = os.path.expanduser("~/.kaggle/kaggle.json")
-    if not os.path.exists(kaggle_json):
-        raise FileNotFoundError(f"未找到 {kaggle_json}，请先在 Kaggle 设置中下载并放置 API Key。")
-
-    # 循环下载
-    for name, kaggle_id in datasets.items():
-        print(f"📥 正在下载 {name} ({kaggle_id}) ...")
-        cmd = [
-            "kaggle", "datasets", "download", "-d", kaggle_id,
-            "-p", os.path.join(save_dir, name), "--unzip"
-        ]
-        subprocess.run(cmd, check=True)
-        print(f"✅ {name} 下载完成\n")
-
-    print("🎉 所有数据集已下载到", save_dir)

lib/logger.py

@@ -4,47 +4,43 @@ from datetime import datetime
 
 
 def get_logger(root, name=None, debug=True):
-    """
+    # when debug is true, show DEBUG and INFO in screen
-    创建带时间戳的日志记录器
+    # when debug is false, show DEBUG in file and info in both screen&file
-    :param root: 日志文件目录
+    # INFO will always be in screen
-    :param name: 日志名称
+    # create a logger
-    :param debug: 是否开启调试模式
-    """
     logger = logging.getLogger(name)
+    # critical > error > warning > info > debug > notset
     logger.setLevel(logging.DEBUG)
 
-    # 避免重复添加 Handler
+    # define the formate
-    if logger.hasHandlers():
+    formatter = logging.Formatter('%(asctime)s: %(message)s', "%m/%d %H:%M")
-        logger.handlers.clear()
+    # create another handler for output log to console
-
-    # 时间格式改为 年/月/日 时:分:秒
-    formatter = logging.Formatter('%(asctime)s - %(message)s', "%Y/%m/%d %H:%M:%S")
-
-    # 控制台输出
     console_handler = logging.StreamHandler()
-    console_handler.setFormatter(formatter)
+    if debug:
-    console_handler.setLevel(logging.DEBUG if debug else logging.INFO)
+        console_handler.setLevel(logging.DEBUG)
-    logger.addHandler(console_handler)
+    else:
-
+        console_handler.setLevel(logging.INFO)
-    # 文件输出（仅非 debug 模式）
+        # create a handler for write log to file
-    if not debug:
-        os.makedirs(root, exist_ok=True)
         logfile = os.path.join(root, 'run.log')
-        print(f"Create Log File in: {logfile}")
+        print('Creat Log File in: ', logfile)
-        file_handler = logging.FileHandler(logfile, mode='w', encoding='utf-8')
+        file_handler = logging.FileHandler(logfile, mode='w')
         file_handler.setLevel(logging.DEBUG)
         file_handler.setFormatter(formatter)
+    console_handler.setFormatter(formatter)
+    # add Handler to logger
+    logger.addHandler(console_handler)
+    if not debug:
         logger.addHandler(file_handler)
-
     return logger
 
 
 if __name__ == '__main__':
-    time_str = datetime.now().strftime('%Y%m%d%H%M%S')
+    time = datetime.now().strftime('%Y%m%d%H%M%S')
-    print(time_str)
+    print(time)
-
+    logger = get_logger('./log.txt', debug=True)
-    logger = get_logger('./logs', debug=False)  # 改成 False 测试文件输出
     logger.debug('this is a {} debug message'.format(1))
     logger.info('this is an info message')
     logger.debug('this is a debug message')
     logger.info('this is an info message')
+    logger.debug('this is a debug message')
+    logger.info('this is an info message')

lib/training_stats.py

@@ -1,72 +0,0 @@
-# 新建 lib/training_stats.py
-import time
-import psutil
-import torch
-import os
-
-
-class TrainingStats:
-    def __init__(self, device):
-        self.device = device
-        self.reset()
-
-    def reset(self):
-        self.gpu_mem_usage_list = []
-        self.cpu_mem_usage_list = []
-        self.train_time_list = []
-        self.infer_time_list = []
-        self.total_iters = 0
-        self.start_time = None
-        self.end_time = None
-
-    def start_training(self):
-        self.start_time = time.time()
-
-    def end_training(self):
-        self.end_time = time.time()
-
-    def record_step_time(self, duration, mode):
-        """记录单步耗时和总迭代次数"""
-        if mode == 'train':
-            self.train_time_list.append(duration)
-        else:
-            self.infer_time_list.append(duration)
-        self.total_iters += 1
-
-    def record_memory_usage(self):
-        """记录当前 GPU 和 CPU 内存占用"""
-        process = psutil.Process(os.getpid())
-        cpu_mem = process.memory_info().rss / (1024 ** 2)
-
-        if torch.cuda.is_available():
-            gpu_mem = torch.cuda.max_memory_allocated(device=self.device) / (1024 ** 2)
-            torch.cuda.reset_peak_memory_stats(device=self.device)
-        else:
-            gpu_mem = 0.0
-
-        self.cpu_mem_usage_list.append(cpu_mem)
-        self.gpu_mem_usage_list.append(gpu_mem)
-
-    def report(self, logger):
-        """在训练结束时输出汇总统计"""
-        if not self.start_time or not self.end_time:
-            logger.warning("TrainingStats: start/end time not recorded properly.")
-            return
-
-        total_time = self.end_time - self.start_time
-        avg_gpu_mem = sum(self.gpu_mem_usage_list) / len(self.gpu_mem_usage_list) if self.gpu_mem_usage_list else 0
-        avg_cpu_mem = sum(self.cpu_mem_usage_list) / len(self.cpu_mem_usage_list) if self.cpu_mem_usage_list else 0
-        avg_train_time = sum(self.train_time_list) / len(self.train_time_list) if self.train_time_list else 0
-        avg_infer_time = sum(self.infer_time_list) / len(self.infer_time_list) if self.infer_time_list else 0
-        iters_per_sec = self.total_iters / total_time if total_time > 0 else 0
-
-        logger.info("===== Training Summary =====")
-        logger.info(f"Total training time: {total_time:.2f} s")
-        logger.info(f"Total iterations: {self.total_iters}")
-        logger.info(f"Average iterations per second: {iters_per_sec:.2f}")
-        logger.info(f"Average GPU Memory Usage: {avg_gpu_mem:.2f} MB")
-        logger.info(f"Average CPU Memory Usage: {avg_cpu_mem:.2f} MB")
-        if avg_train_time:
-            logger.info(f"Average training step time: {avg_train_time*1000:.2f} ms")
-        if avg_infer_time:
-            logger.info(f"Average inference step time: {avg_infer_time*1000:.2f} ms")

run.py

@@ -17,9 +17,6 @@ from dataloader.loader_selector import get_dataloader
 from trainer.trainer_selector import select_trainer
 import yaml
 
-
-
-
 def main():
     args = parse_args()
 

trainer/Trainer.py

@@ -2,7 +2,6 @@ import math
 import os
 import time
 import copy
-import psutil
 from tqdm import tqdm
 
 import torch
@@ -10,73 +9,6 @@ from lib.logger import get_logger
 from lib.loss_function import all_metrics
 
 
-class TrainingStats:
-    def __init__(self, device):
-        self.device = device
-        self.reset()
-
-    def reset(self):
-        self.gpu_mem_usage_list = []
-        self.cpu_mem_usage_list = []
-        self.train_time_list = []
-        self.infer_time_list = []
-        self.total_iters = 0
-        self.start_time = None
-        self.end_time = None
-
-    def start_training(self):
-        self.start_time = time.time()
-
-    def end_training(self):
-        self.end_time = time.time()
-
-    def record_step_time(self, duration, mode):
-        """记录单步耗时和总迭代次数"""
-        if mode == 'train':
-            self.train_time_list.append(duration)
-        else:
-            self.infer_time_list.append(duration)
-        self.total_iters += 1
-
-    def record_memory_usage(self):
-        """记录当前 GPU 和 CPU 内存占用"""
-        process = psutil.Process(os.getpid())
-        cpu_mem = process.memory_info().rss / (1024 ** 2)
-
-        if torch.cuda.is_available():
-            gpu_mem = torch.cuda.max_memory_allocated(device=self.device) / (1024 ** 2)
-            torch.cuda.reset_peak_memory_stats(device=self.device)
-        else:
-            gpu_mem = 0.0
-
-        self.cpu_mem_usage_list.append(cpu_mem)
-        self.gpu_mem_usage_list.append(gpu_mem)
-
-    def report(self, logger):
-        """在训练结束时输出汇总统计"""
-        if not self.start_time or not self.end_time:
-            logger.warning("TrainingStats: start/end time not recorded properly.")
-            return
-
-        total_time = self.end_time - self.start_time
-        avg_gpu_mem = sum(self.gpu_mem_usage_list) / len(self.gpu_mem_usage_list) if self.gpu_mem_usage_list else 0
-        avg_cpu_mem = sum(self.cpu_mem_usage_list) / len(self.cpu_mem_usage_list) if self.cpu_mem_usage_list else 0
-        avg_train_time = sum(self.train_time_list) / len(self.train_time_list) if self.train_time_list else 0
-        avg_infer_time = sum(self.infer_time_list) / len(self.infer_time_list) if self.infer_time_list else 0
-        iters_per_sec = self.total_iters / total_time if total_time > 0 else 0
-
-        logger.info("===== Training Summary =====")
-        logger.info(f"Total training time: {total_time:.2f} s")
-        logger.info(f"Total iterations: {self.total_iters}")
-        logger.info(f"Average iterations per second: {iters_per_sec:.2f}")
-        logger.info(f"Average GPU Memory Usage: {avg_gpu_mem:.2f} MB")
-        logger.info(f"Average CPU Memory Usage: {avg_cpu_mem:.2f} MB")
-        if avg_train_time:
-            logger.info(f"Average training step time: {avg_train_time*1000:.2f} ms")
-        if avg_infer_time:
-            logger.info(f"Average inference step time: {avg_infer_time*1000:.2f} ms")
-
-
 class Trainer:
     def __init__(self, model, loss, optimizer, train_loader, val_loader, test_loader,
                  scaler, args, lr_scheduler=None):
@@ -103,9 +35,6 @@ class Trainer:
         self.logger = get_logger(args['log_dir'], name=self.model.__class__.__name__, debug=args['debug'])
         self.logger.info(f"Experiment log path in: {args['log_dir']}")
 
-        # Stats tracker
-        self.stats = TrainingStats(device=args['device'])
-
     def _run_epoch(self, epoch, dataloader, mode):
         if mode == 'train':
             self.model.train()
@@ -120,8 +49,6 @@ class Trainer:
         with torch.set_grad_enabled(optimizer_step):
             with tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar:
                 for batch_idx, (data, target) in enumerate(dataloader):
-                    start_time = time.time()
-
                     label = target[..., :self.args['output_dim']]
                     output = self.model(data).to(self.args['device'])
 
@@ -137,25 +64,19 @@ class Trainer:
                             torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm'])
                         self.optimizer.step()
 
-                    step_time = time.time() - start_time
-                    self.stats.record_step_time(step_time, mode)
-
                     total_loss += loss.item()
 
                     if mode == 'train' and (batch_idx + 1) % self.args['log_step'] == 0:
                         self.logger.info(
                             f'Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}')
 
+                    # 更新 tqdm 的进度
                     pbar.update(1)
                     pbar.set_postfix(loss=loss.item())
 
         avg_loss = total_loss / len(dataloader)
         self.logger.info(
             f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s')
-
-        # 记录内存
-        self.stats.record_memory_usage()
-
         return avg_loss
 
     def train_epoch(self, epoch):
@@ -172,9 +93,7 @@ class Trainer:
         best_loss, best_test_loss = float('inf'), float('inf')
         not_improved_count = 0
 
-        self.stats.start_training()
         self.logger.info("Training process started")
-
         for epoch in range(1, self.args['epochs'] + 1):
             train_epoch_loss = self.train_epoch(epoch)
             val_epoch_loss = self.val_epoch(epoch)
@@ -188,6 +107,7 @@ class Trainer:
                 best_loss = val_epoch_loss
                 not_improved_count = 0
                 best_model = copy.deepcopy(self.model.state_dict())
+                torch.save(best_model, self.best_path)
                 self.logger.info('Best validation model saved!')
             else:
                 not_improved_count += 1
@@ -199,6 +119,7 @@ class Trainer:
 
             if test_epoch_loss < best_test_loss:
                 best_test_loss = test_epoch_loss
+                torch.save(best_test_model, self.best_test_path)
                 best_test_model = copy.deepcopy(self.model.state_dict())
 
         if not self.args['debug']:
@@ -206,9 +127,6 @@ class Trainer:
             torch.save(best_test_model, self.best_test_path)
             self.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}")
 
-        self.stats.end_training()
-        self.stats.report(self.logger)
-
         self._finalize_training(best_model, best_test_model)
 
     def _finalize_training(self, best_model, best_test_model):
@@ -243,6 +161,10 @@ class Trainer:
             y_pred = torch.cat(y_pred, dim=0)
         y_true = torch.cat(y_true, dim=0)
 
+        # 你在这里需要把y_pred和y_true保存下来
+        # torch.save(y_pred, "./test/PEMS07/y_pred_D.pt") # [3566,12,170,1]
+        # torch.save(y_true, "./test/PEMSD8/y_true.pt") # [3566,12,170,1]
+
         for t in range(y_true.shape[1]):
             mae, rmse, mape = all_metrics(y_pred[:, t, ...], y_true[:, t, ...],
                                           args['mae_thresh'], args['mape_thresh'])

trainer/Trainer_old.py

@@ -1,176 +0,0 @@
-import math
-import os
-import time
-import copy
-from tqdm import tqdm
-
-import torch
-from lib.logger import get_logger
-from lib.loss_function import all_metrics
-
-
-class Trainer:
-    def __init__(self, model, loss, optimizer, train_loader, val_loader, test_loader,
-                 scaler, args, lr_scheduler=None):
-        self.model = model
-        self.loss = loss
-        self.optimizer = optimizer
-        self.train_loader = train_loader
-        self.val_loader = val_loader
-        self.test_loader = test_loader
-        self.scaler = scaler
-        self.args = args
-        self.lr_scheduler = lr_scheduler
-        self.train_per_epoch = len(train_loader)
-        self.val_per_epoch = len(val_loader) if val_loader else 0
-
-        # Paths for saving models and logs
-        self.best_path = os.path.join(args['log_dir'], 'best_model.pth')
-        self.best_test_path = os.path.join(args['log_dir'], 'best_test_model.pth')
-        self.loss_figure_path = os.path.join(args['log_dir'], 'loss.png')
-
-        # Initialize logger
-        if not os.path.isdir(args['log_dir']) and not args['debug']:
-            os.makedirs(args['log_dir'], exist_ok=True)
-        self.logger = get_logger(args['log_dir'], name=self.model.__class__.__name__, debug=args['debug'])
-        self.logger.info(f"Experiment log path in: {args['log_dir']}")
-
-    def _run_epoch(self, epoch, dataloader, mode):
-        if mode == 'train':
-            self.model.train()
-            optimizer_step = True
-        else:
-            self.model.eval()
-            optimizer_step = False
-
-        total_loss = 0
-        epoch_time = time.time()
-
-        with torch.set_grad_enabled(optimizer_step):
-            with tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar:
-                for batch_idx, (data, target) in enumerate(dataloader):
-                    label = target[..., :self.args['output_dim']]
-                    output = self.model(data).to(self.args['device'])
-
-                    if self.args['real_value']:
-                        output = self.scaler.inverse_transform(output)
-
-                    loss = self.loss(output, label)
-                    if optimizer_step and self.optimizer is not None:
-                        self.optimizer.zero_grad()
-                        loss.backward()
-
-                        if self.args['grad_norm']:
-                            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm'])
-                        self.optimizer.step()
-
-                    total_loss += loss.item()
-
-                    if mode == 'train' and (batch_idx + 1) % self.args['log_step'] == 0:
-                        self.logger.info(
-                            f'Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}')
-
-                    # 更新 tqdm 的进度
-                    pbar.update(1)
-                    pbar.set_postfix(loss=loss.item())
-
-        avg_loss = total_loss / len(dataloader)
-        self.logger.info(
-            f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s')
-        return avg_loss
-
-    def train_epoch(self, epoch):
-        return self._run_epoch(epoch, self.train_loader, 'train')
-
-    def val_epoch(self, epoch):
-        return self._run_epoch(epoch, self.val_loader or self.test_loader, 'val')
-
-    def test_epoch(self, epoch):
-        return self._run_epoch(epoch, self.test_loader, 'test')
-
-    def train(self):
-        best_model, best_test_model = None, None
-        best_loss, best_test_loss = float('inf'), float('inf')
-        not_improved_count = 0
-
-        self.logger.info("Training process started")
-        for epoch in range(1, self.args['epochs'] + 1):
-            train_epoch_loss = self.train_epoch(epoch)
-            val_epoch_loss = self.val_epoch(epoch)
-            test_epoch_loss = self.test_epoch(epoch)
-
-            if train_epoch_loss > 1e6:
-                self.logger.warning('Gradient explosion detected. Ending...')
-                break
-
-            if val_epoch_loss < best_loss:
-                best_loss = val_epoch_loss
-                not_improved_count = 0
-                best_model = copy.deepcopy(self.model.state_dict())
-                self.logger.info('Best validation model saved!')
-            else:
-                not_improved_count += 1
-
-            if self.args['early_stop'] and not_improved_count == self.args['early_stop_patience']:
-                self.logger.info(
-                    f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops.")
-                break
-
-            if test_epoch_loss < best_test_loss:
-                best_test_loss = test_epoch_loss
-                best_test_model = copy.deepcopy(self.model.state_dict())
-
-        if not self.args['debug']:
-            torch.save(best_model, self.best_path)
-            torch.save(best_test_model, self.best_test_path)
-            self.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}")
-
-        self._finalize_training(best_model, best_test_model)
-
-    def _finalize_training(self, best_model, best_test_model):
-        self.model.load_state_dict(best_model)
-        self.logger.info("Testing on best validation model")
-        self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)
-
-        self.model.load_state_dict(best_test_model)
-        self.logger.info("Testing on best test model")
-        self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)
-
-    @staticmethod
-    def test(model, args, data_loader, scaler, logger, path=None):
-        if path:
-            checkpoint = torch.load(path)
-            model.load_state_dict(checkpoint['state_dict'])
-            model.to(args['device'])
-
-        model.eval()
-        y_pred, y_true = [], []
-
-        with torch.no_grad():
-            for data, target in data_loader:
-                label = target[..., :args['output_dim']]
-                output = model(data)
-                y_pred.append(output)
-                y_true.append(label)
-
-        if args['real_value']:
-            y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
-        else:
-            y_pred = torch.cat(y_pred, dim=0)
-        y_true = torch.cat(y_true, dim=0)
-
-        # 你在这里需要把y_pred和y_true保存下来
-        # torch.save(y_pred, "./test/PEMS07/y_pred_D.pt") # [3566,12,170,1]
-        # torch.save(y_true, "./test/PEMS08/y_true.pt") # [3566,12,170,1]
-
-        for t in range(y_true.shape[1]):
-            mae, rmse, mape = all_metrics(y_pred[:, t, ...], y_true[:, t, ...],
-                                          args['mae_thresh'], args['mape_thresh'])
-            logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-
-        mae, rmse, mape = all_metrics(y_pred, y_true, args['mae_thresh'], args['mape_thresh'])
-        logger.info(f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-
-    @staticmethod
-    def _compute_sampling_threshold(global_step, k):
-        return k / (k + math.exp(global_step / k))