DualFlow/models/causal-inference/data_trainer.py

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import os
from sklearn.metrics import mean_absolute_error, mean_squared_error
import logging

class DataTrainer:
    def __init__(self, model, args, preprocessor, optimizer=None, scheduler=None, logger=None):
        self.model = model
        self.args = args
        self.preprocessor = preprocessor
        self.device = args.device
        self.logger = logger if logger is not None else self._default_logger()
        
        self.criterion = nn.MSELoss()
        self.optimizer = optimizer if optimizer is not None else optim.Adam(
            model.parameters(),
            lr=args.lr,
            weight_decay=args.weight_decay
        )
        self.scheduler = scheduler
        
        self.train_losses = []
        self.val_losses = []
        self.train_mae = []
        self.val_mae = []
        self.best_val_loss = float('inf')
        self.early_stop_counter = 0
        self.model_save_path = os.path.join('models', 'best_model.pth')
        self.final_model_path = os.path.join('models', 'final_model.pth')
        
        if not os.path.exists('models'):
            os.makedirs('models')
        if not os.path.exists('plots'):
            os.makedirs('plots')
        
    def _default_logger(self):
        logger = logging.getLogger('DefaultLogger')
        logger.setLevel(logging.INFO)
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        console_handler.setFormatter(formatter)
        if not logger.handlers:
            logger.addHandler(console_handler)
        return logger
    
    @torch.enable_grad()
    def train_epoch(self, train_loader, adj, max_batches=None):
        """
        支持通过 max_batches 限制本 epoch 使用的批次数（用于 RL 快速评估）。
        """
        self.model.train()
        total_loss = 0.0
        all_outputs = []
        all_targets = []
        adj = adj.to(self.device)

        for batch_idx, (data, target) in enumerate(train_loader):
            if max_batches is not None and batch_idx >= max_batches:
                break
            data, target = data.to(self.device), target.to(self.device)
            data = data.unsqueeze(-1)  # (batch_size, 145, 1)
            
            self.optimizer.zero_grad()
            output = self.model(data, adj)  # (batch_size, 145, 47)
            output = output.mean(dim=1)     # (batch_size, 47)
            
            loss = self.criterion(output, target)
            loss.backward()
            if self.args.grad_clip > 0:
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.grad_clip)
            self.optimizer.step()
            
            total_loss += loss.item()
            all_outputs.append(output.detach().cpu().numpy())
            all_targets.append(target.detach().cpu().numpy())
        
        if len(all_outputs) == 0:
            # 在极小数据/极小 batch 情况下的保护
            return float('inf'), float('inf')
        
        all_outputs = np.vstack(all_outputs)
        all_targets = np.vstack(all_targets)
        mae = mean_absolute_error(all_targets, all_outputs)
        
        avg_loss = total_loss / (min(len(train_loader), max_batches) if max_batches else len(train_loader))
        self.train_losses.append(avg_loss)
        self.train_mae.append(mae)
        return avg_loss, mae
    
    @torch.no_grad()
    def validate(self, val_loader, adj, max_batches=None):
        """
        支持通过 max_batches 限制验证批次数（用于 RL 快速评估）。
        """
        self.model.eval()
        total_loss = 0.0
        all_outputs = []
        all_targets = []
        adj = adj.to(self.device)
        
        for batch_idx, (data, target) in enumerate(val_loader):
            if max_batches is not None and batch_idx >= max_batches:
                break
            data, target = data.to(self.device), target.to(self.device)
            data = data.unsqueeze(-1)  # (batch_size, 145, 1)
            output = self.model(data, adj)  # (batch_size, 145, 47)
            output = output.mean(dim=1)     # (batch_size, 47)
            loss = self.criterion(output, target)
            total_loss += loss.item()
            all_outputs.append(output.cpu().numpy())
            all_targets.append(target.cpu().numpy())
        
        if len(all_outputs) == 0:
            return float('inf'), float('inf')
        
        all_outputs = np.vstack(all_outputs)
        all_targets = np.vstack(all_targets)
        mae = mean_absolute_error(all_targets, all_outputs)
        avg_loss = total_loss / (min(len(val_loader), max_batches) if max_batches else len(val_loader))
        self.val_losses.append(avg_loss)
        self.val_mae.append(mae)
        return avg_loss, mae
    
    def train(self, train_loader, val_loader, adj, epochs=None):
        adj = adj.to(self.device)
        epochs = epochs if epochs is not None else self.args.epochs
        self.logger.info(f"开始训练，共 {epochs} 个epoch")
        
        for epoch in tqdm(range(epochs)):
            train_loss, train_mae = self.train_epoch(train_loader, adj)
            val_loss, val_mae = self.validate(val_loader, adj)
            if self.scheduler is not None:
                self.scheduler.step(val_loss)
            
            # 保存最佳模型
            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': self.model.state_dict(),
                    'optimizer_state_dict': self.optimizer.state_dict(),
                    'loss': val_loss,
                    'args': self.args  # 保存训练参数
                }, self.model_save_path)
                self.early_stop_counter = 0
                self.logger.info(f"已更新最佳模型到 {self.model_save_path}")  # 日志
            else:
                self.early_stop_counter += 1
                if self.early_stop_counter >= self.args.patience:
                    self.logger.info(f"早停机制触发，在第 {epoch+1} 轮停止训练")
                    break
            
            if (epoch + 1) % 10 == 0:
                self.logger.info(f'Epoch {epoch+1}/{epochs}, '
                                 f'Train Loss: {train_loss:.6f}, Train MAE: {train_mae:.6f}, '
                                 f'Val Loss: {val_loss:.6f}, Val MAE: {val_mae:.6f}')
        
        self.plot_losses()
        self.plot_mae()
        
        # 加载最佳模型
        checkpoint = torch.load(self.model_save_path, map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.logger.info(f"加载最佳模型（第 {checkpoint['epoch']+1} 轮，验证损失: {checkpoint['loss']:.6f}）")
        
        # 保存最终训练完成的模型（加载最佳模型后）
        torch.save({
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'best_val_loss': self.best_val_loss,
            'args': self.args
        }, self.final_model_path)
        self.logger.info(f"已保存最终模型到 {self.final_model_path}")
        
        return self.model
    
    @torch.no_grad()
    def test(self, test_loader, adj):
        self.model.eval()
        total_loss = 0.0
        all_outputs = []
        all_targets = []
        adj = adj.to(self.device)
        
        for data, target in test_loader:
            data, target = data.to(self.device), target.to(self.device)
            data = data.unsqueeze(-1)
            output = self.model(data, adj)
            output = output.mean(dim=1)
            loss = self.criterion(output, target)
            total_loss += loss.item()
            all_outputs.append(output.cpu().numpy())
            all_targets.append(target.cpu().numpy())
        
        all_outputs = np.vstack(all_outputs)
        all_targets = np.vstack(all_targets)
        mse = total_loss / len(test_loader)
        mae = mean_absolute_error(all_targets, all_outputs)
        rmse = np.sqrt(mean_squared_error(all_targets, all_outputs))
        
        all_outputs_original = self.preprocessor.inverse_transform_targets(all_outputs)
        all_targets_original = self.preprocessor.inverse_transform_targets(all_targets)
        original_mse = mean_squared_error(all_targets_original, all_outputs_original)
        original_mae = mean_absolute_error(all_targets_original, all_outputs_original)
        original_rmse = np.sqrt(original_mse)
        
        self.logger.info(f'Test Loss (normalized): MSE={mse:.6f}, MAE={mae:.6f}, RMSE={rmse:.6f}')
        self.logger.info(f'Test Loss (original scale): MSE={original_mse:.6f}, MAE={original_mae:.6f}, RMSE={original_rmse:.6f}')
        self.plot_predictions(all_outputs_original, all_targets_original)
        return {
            'normalized_mse': mse,
            'normalized_mae': mae,
            'normalized_rmse': rmse,
            'original_mse': original_mse,
            'original_mae': original_mae,
            'original_rmse': original_rmse,
            'predictions': all_outputs_original,
            'targets': all_targets_original
        }
    
    def plot_losses(self):
        plt.figure(figsize=(10, 6))
        plt.plot(self.train_losses, label='Train Loss')
        plt.plot(self.val_losses, label='Validation Loss')
        plt.xlabel('Epoch')
        plt.ylabel('MSE Loss')
        plt.title('Training and Validation Loss')
        plt.legend()
        plt.savefig('plots/loss_curve.png')
        plt.close()
    
    def plot_mae(self):
        plt.figure(figsize=(10, 6))
        plt.plot(self.train_mae, label='Train MAE')
        plt.plot(self.val_mae, label='Validation MAE')
        plt.xlabel('Epoch')
        plt.ylabel('MAE')
        plt.title('Training and Validation MAE')
        plt.legend()
        plt.savefig('plots/mae_curve.png')
        plt.close()
    
    def plot_predictions(self, predictions, targets):
        num_plots = min(3, self.args.num_targets)
        plt.figure(figsize=(15, 5*num_plots))
        for i in range(num_plots):
            plt.subplot(num_plots, 1, i+1)
            plt.plot(targets[:100, i], label='True Value')
            plt.plot(predictions[:100, i], label='Predicted Value')
            plt.xlabel('Time Step')
            plt.ylabel(f'Target {i+1}')
            plt.title(f'Prediction vs True Value for Target {i+1}')
            plt.legend()
        plt.tight_layout()
        plt.savefig('plots/prediction_examples.png')
        plt.close()