src/models/pad_predictor.py

"""
PAD预测器模型
PAD Predictor Model for emotion and physiological state change prediction

该模型实现了一个多层感知机(MLP)来预测用户情绪和生理状态的变化。
输入：7维 (User PAD 3维 + Vitality 1维 + Current PAD 3维)
输出：5维 (ΔPAD 3维 + ΔPressure 1维 + Confidence 1维)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Any, Optional, Tuple
import logging


class PADPredictor(nn.Module):
    """
    PAD情绪和生理状态变化预测器
    
    网络架构：
    - 输入层：7维特征
    - 隐藏层1：128神经元 + ReLU + Dropout(0.3)
    - 隐藏层2：64神经元 + ReLU + Dropout(0.3)
    - 隐藏层3：32神经元 + ReLU
    - 输出层：5神经元 + Linear激活
    """
    
    def __init__(self,
                 input_dim: int = 10,
                 output_dim: int = 4,
                 hidden_dims: list = [128, 64, 32],
                 dropout_rate: float = 0.3,
                 weight_init: str = "xavier_uniform",
                 bias_init: str = "zeros"):
        """
        初始化PAD预测器

        Args:
            input_dim: 输入维度，默认10维（7原始特征+3差异特征）
            output_dim: 输出维度，默认4维（ΔPAD 3维 + ΔPressure 1维）
            hidden_dims: 隐藏层维度列表
            dropout_rate: Dropout概率
            weight_init: 权重初始化方法
            bias_init: 偏置初始化方法
        """
        super(PADPredictor, self).__init__()
        
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.hidden_dims = hidden_dims
        self.dropout_rate = dropout_rate
        self.weight_init = weight_init
        self.bias_init = bias_init
        
        # 构建网络层
        self._build_network()
        
        # 初始化权重
        self._initialize_weights()
        
        # 设置日志
        self.logger = logging.getLogger(__name__)
        
    def _build_network(self):
        """构建网络架构"""
        layers = []
        
        # 输入层到第一个隐藏层
        layers.append(nn.Linear(self.input_dim, self.hidden_dims[0]))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(self.dropout_rate))
        
        # 中间隐藏层
        for i in range(len(self.hidden_dims) - 1):
            layers.append(nn.Linear(self.hidden_dims[i], self.hidden_dims[i + 1]))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(self.dropout_rate))
        
        # 最后一个隐藏层（不加Dropout）
        layers.append(nn.Linear(self.hidden_dims[-1], self.output_dim))
        # 输出层使用线性激活（回归任务）
        
        # 构建序列网络
        self.network = nn.Sequential(*layers)
        
    def _initialize_weights(self):
        """初始化网络权重"""
        for m in self.modules():
            if isinstance(m, nn.Linear):
                # 权重初始化
                if self.weight_init == "xavier_uniform":
                    nn.init.xavier_uniform_(m.weight)
                elif self.weight_init == "xavier_normal":
                    nn.init.xavier_normal_(m.weight)
                elif self.weight_init == "kaiming_uniform":
                    nn.init.kaiming_uniform_(m.weight, nonlinearity='relu')
                elif self.weight_init == "kaiming_normal":
                    nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                else:
                    # 默认使用Xavier均匀初始化
                    nn.init.xavier_uniform_(m.weight)
                
                # 偏置初始化
                if self.bias_init == "zeros":
                    nn.init.zeros_(m.bias)
                elif self.bias_init == "ones":
                    nn.init.ones_(m.bias)
                else:
                    # 默认使用零初始化
                    nn.init.zeros_(m.bias)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播
        
        Args:
            x: 输入张量，形状为 (batch_size, input_dim)
            
        Returns:
            输出张量，形状为 (batch_size, output_dim)
        """
        # 输入验证
        if x.dim() != 2:
            raise ValueError(f"输入张量应该是2维的 (batch_size, input_dim)，但得到的是 {x.dim()}维")
        
        if x.size(1) != self.input_dim:
            raise ValueError(f"输入维度应该是 {self.input_dim}，但得到的是 {x.size(1)}")
        
        # 前向传播
        output = self.network(x)
        
        return output
    
    def predict_components(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        预测并分解输出组件
        
        Args:
            x: 输入张量
            
        Returns:
            包含各组件的字典：
            - 'delta_pad': ΔPAD (3维)
            - 'delta_pressure': ΔPressure (1维)
            - 'confidence': Confidence (1维)
        """
        output = self.forward(x)
        
        # 分解输出
        delta_pad = output[:, :3]  # 前3维：ΔPAD
        delta_pressure = output[:, 3:4]  # 第4维：ΔPressure
        confidence = output[:, 4:5]  # 第5维：Confidence
        
        return {
            'delta_pad': delta_pad,
            'delta_pressure': delta_pressure,
            'confidence': confidence
        }
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        
        Returns:
            包含模型信息的字典
        """
        total_params = sum(p.numel() for p in self.parameters())
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
        
        return {
            'model_type': 'PADPredictor',
            'input_dim': self.input_dim,
            'output_dim': self.output_dim,
            'hidden_dims': self.hidden_dims,
            'dropout_rate': self.dropout_rate,
            'weight_init': self.weight_init,
            'bias_init': self.bias_init,
            'total_parameters': total_params,
            'trainable_parameters': trainable_params,
            'architecture': str(self.network)
        }
    
    def save_model(self, filepath: str, include_optimizer: bool = False, optimizer: Optional[torch.optim.Optimizer] = None):
        """
        保存模型
        
        Args:
            filepath: 保存路径
            include_optimizer: 是否包含优化器状态
            optimizer: 优化器对象
        """
        save_dict = {
            'model_state_dict': self.state_dict(),
            'model_config': {
                'input_dim': self.input_dim,
                'output_dim': self.output_dim,
                'hidden_dims': self.hidden_dims,
                'dropout_rate': self.dropout_rate,
                'weight_init': self.weight_init,
                'bias_init': self.bias_init
            },
            'model_info': self.get_model_info()
        }
        
        if include_optimizer and optimizer is not None:
            save_dict['optimizer_state_dict'] = optimizer.state_dict()
        
        torch.save(save_dict, filepath)
        self.logger.info(f"模型已保存到: {filepath}")
    
    @classmethod
    def load_model(cls, filepath: str, device: str = 'cpu'):
        """
        加载模型
        
        Args:
            filepath: 模型文件路径
            device: 设备类型
            
        Returns:
            PADPredictor实例
        """
        # 针对 PyTorch 2.6+ 的安全性更新，显式设置 weights_only=False 以加载包含 numpy 标量的 checkpoint
        try:
            checkpoint = torch.load(filepath, map_location=device, weights_only=False)
        except TypeError:
            # 兼容旧版本 PyTorch
            checkpoint = torch.load(filepath, map_location=device)
        
        # 尝试获取配置，兼容不同的键名
        model_config = checkpoint.get('model_config') or checkpoint.get('config')
        if model_config is None:
            raise KeyError("Checkpoint 中未找到 'model_config' 或 'config' 键")
            
        model = cls(**model_config)
        
        # 加载权重
        model.load_state_dict(checkpoint['model_state_dict'])
        
        logging.info(f"模型已从 {filepath} 加载")
        return model
    
    def freeze_layers(self, layer_names: list = None):
        """
        冻结指定层
        
        Args:
            layer_names: 要冻结的层名称列表，如果为None则冻结所有层
        """
        if layer_names is None:
            # 冻结所有参数
            for param in self.parameters():
                param.requires_grad = False
            self.logger.info("所有层已被冻结")
        else:
            # 冻结指定层
            for name, param in self.named_parameters():
                for layer_name in layer_names:
                    if layer_name in name:
                        param.requires_grad = False
                        break
            self.logger.info(f"指定层 {layer_names} 已被冻结")
    
    def unfreeze_layers(self, layer_names: list = None):
        """
        解冻指定层
        
        Args:
            layer_names: 要解冻的层名称列表，如果为None则解冻所有层
        """
        if layer_names is None:
            # 解冻所有参数
            for param in self.parameters():
                param.requires_grad = True
            self.logger.info("所有层已被解冻")
        else:
            # 解冻指定层
            for name, param in self.named_parameters():
                for layer_name in layer_names:
                    if layer_name in name:
                        param.requires_grad = True
                        break
            self.logger.info(f"指定层 {layer_names} 已被解冻")


def create_pad_predictor(config: Optional[Dict[str, Any]] = None) -> PADPredictor:
    """
    创建PAD预测器的工厂函数
    
    Args:
        config: 配置字典，包含模型参数
        
    Returns:
        PADPredictor实例
    """
    if config is None:
        # 使用默认配置
        return PADPredictor()
    
    # 从配置中提取参数
    model_config = config.get('architecture', {})
    
    return PADPredictor(
        input_dim=config.get('dimensions', {}).get('input_dim', 10),  # 默认10维（7原始+3差异）
        output_dim=config.get('dimensions', {}).get('output_dim', 4),  # 默认4维（移除confidence）
        hidden_dims=[layer['size'] for layer in model_config.get('hidden_layers', [])],
        dropout_rate=model_config.get('dropout_config', {}).get('rate', 0.3),
        weight_init=config.get('initialization', {}).get('weight_init', 'xavier_uniform'),
        bias_init=config.get('initialization', {}).get('bias_init', 'zeros')
    )


if __name__ == "__main__":
    # 测试代码
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 创建模型
    model = PADPredictor().to(device)
    
    # 打印模型信息
    print("模型信息:")
    info = model.get_model_info()
    for key, value in info.items():
        print(f"  {key}: {value}")
    
    # 测试前向传播
    batch_size = 4
    x = torch.randn(batch_size, 7).to(device)
    
    with torch.no_grad():
        output = model(x)
        components = model.predict_components(x)
        
        print(f"\n输入形状: {x.shape}")
        print(f"输出形状: {output.shape}")
        print(f"ΔPAD形状: {components['delta_pad'].shape}")
        print(f"ΔPressure形状: {components['delta_pressure'].shape}")
        print(f"Confidence形状: {components['confidence'].shape}")
    
    print("\n模型测试完成！")