""" TFT (Temporal Fusion Transformer) 模型实现 简化版本，支持静态特征、已知未来特征和观测时序特征 """ import torch import torch.nn as nn import torch.nn.functional as F import math class TFTEncoder(nn.Module): """TFT编码器""" def __init__(self, d_model, nhead, num_layers, dim_feedforward, dropout=0.1): super(TFTEncoder, self).__init__() encoder_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True ) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers) def forward(self, x, mask=None): """ 前向传播 参数: x: 输入 [batch_size, seq_len, d_model] mask: 注意力mask [batch_size, seq_len] 返回: output: 编码输出 [batch_size, seq_len, d_model] """ return self.encoder(x, src_key_padding_mask=mask) class TFTDecoder(nn.Module): """TFT解码器""" def __init__(self, d_model, nhead, num_layers, dim_feedforward, dropout=0.1): super(TFTDecoder, self).__init__() decoder_layer = nn.TransformerDecoderLayer( d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True ) self.decoder = nn.TransformerDecoder(decoder_layer, num_layers) def forward(self, tgt, memory, tgt_mask=None, memory_mask=None): """ 前向传播 参数: tgt: 目标序列 [batch_size, tgt_len, d_model] memory: 编码器输出 [batch_size, seq_len, d_model] tgt_mask: 目标mask memory_mask: 记忆mask 返回: output: 解码输出 [batch_size, tgt_len, d_model] """ return self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_key_padding_mask=memory_mask) class TemporalFusionTransformer(nn.Module): """ Temporal Fusion Transformer (简化版) 支持静态特征、已知未来特征和观测时序特征 """ def __init__(self, num_observed_features, num_static_features, num_known_future_features, num_output_features=None, hidden_size=128, num_heads=4, num_encoder_layers=3, num_decoder_layers=3, dim_feedforward=512, dropout=0.1): """ 初始化TFT 参数: num_observed_features: 观测时序特征维度（来自Phased LSTM的输出） num_static_features: 静态特征维度 num_known_future_features: 已知未来特征维度 hidden_size: 隐藏层维度 num_heads: 注意力头数 num_encoder_layers: 编码器层数 num_decoder_layers: 解码器层数 dim_feedforward: 前馈网络维度 dropout: Dropout率 """ super(TemporalFusionTransformer, self).__init__() self.hidden_size = hidden_size self.num_observed_features = num_observed_features self.num_static_features = num_static_features self.num_known_future_features = num_known_future_features self.num_output_features = num_output_features or num_observed_features # 默认输出与观测特征相同 # 特征嵌入 self.observed_embedding = nn.Linear(num_observed_features, hidden_size) self.static_embedding = nn.Linear(num_static_features, hidden_size) self.known_future_embedding = nn.Linear(num_known_future_features, hidden_size) # 位置编码（可选） self.pos_encoder = PositionalEncoding(hidden_size, dropout) # 编码器 self.encoder = TFTEncoder( d_model=hidden_size, nhead=num_heads, num_layers=num_encoder_layers, dim_feedforward=dim_feedforward, dropout=dropout ) # 解码器 self.decoder = TFTDecoder( d_model=hidden_size, nhead=num_heads, num_layers=num_decoder_layers, dim_feedforward=dim_feedforward, dropout=dropout ) # 输出层（预测原始特征） self.output_layer = nn.Linear(hidden_size, self.num_output_features) # 静态特征融合（将静态特征广播到每个时间步） self.static_fusion = nn.Linear(hidden_size * 2, hidden_size) def forward(self, observed_features, static_features, known_future_features, mask=None, debug=False): """ 前向传播 参数: observed_features: 观测时序特征 [batch_size, seq_len, num_observed_features] static_features: 静态特征 [batch_size, num_static_features] known_future_features: 已知未来特征 [batch_size, pred_len, num_known_future_features] mask: 注意力mask [batch_size, seq_len, num_features] 或 [batch_size, seq_len] 如果是3D：1表示有效，0表示缺失；如果某个时间步的所有特征都缺失，则该时间步被mask 如果是2D：True表示需要mask的位置（padding），False表示有效位置 debug: 是否启用调试日志 返回: predictions: 预测值 [batch_size, pred_len, num_observed_features] """ batch_size, seq_len, _ = observed_features.size() pred_len = known_future_features.size(1) if debug: print(f" [TFT] 输入检查:") print(f" observed_features: shape={observed_features.shape}, has_nan={torch.isnan(observed_features).any().item()}") print(f" static_features: shape={static_features.shape}, has_nan={torch.isnan(static_features).any().item()}") print(f" known_future_features: shape={known_future_features.shape}, has_nan={torch.isnan(known_future_features).any().item()}") # 嵌入观测特征 observed_emb = self.observed_embedding(observed_features) # [batch_size, seq_len, hidden_size] if debug and torch.isnan(observed_emb).any(): print(f" ⚠️ observed_emb包含NaN（在embedding后）") observed_emb = self.pos_encoder(observed_emb) if debug and torch.isnan(observed_emb).any(): print(f" ⚠️ observed_emb包含NaN（在pos_encoder后）") # 嵌入静态特征并广播 static_emb = self.static_embedding(static_features) # [batch_size, hidden_size] if debug and torch.isnan(static_emb).any(): print(f" ⚠️ static_emb包含NaN（在embedding后）") static_emb = static_emb.unsqueeze(1).expand(-1, seq_len, -1) # [batch_size, seq_len, hidden_size] # 融合静态特征和观测特征 encoder_input = torch.cat([observed_emb, static_emb], dim=-1) # [batch_size, seq_len, hidden_size*2] if debug and torch.isnan(encoder_input).any(): print(f" ⚠️ encoder_input包含NaN（在concat后）") encoder_input = self.static_fusion(encoder_input) # [batch_size, seq_len, hidden_size] if debug and torch.isnan(encoder_input).any(): print(f" ⚠️ encoder_input包含NaN（在static_fusion后）") # 将3D mask转换为2D mask（如果mask是3D的） # mask: [batch_size, seq_len, num_features] -> [batch_size, seq_len] # 如果某个时间步的所有特征都缺失，则该时间步被mask if mask is not None and mask.dim() == 3: # mask值为1表示有效，0表示缺失 # 如果某个时间步的所有特征都缺失（sum=0），则该时间步应该被mask（True） mask_2d = (mask.sum(dim=-1) == 0).bool() # [batch_size, seq_len] # True表示需要mask的位置（padding），False表示有效位置 # 检查：如果整个batch的所有序列都被mask，则设为None（避免Transformer错误） if mask_2d.all(): mask_2d = None # 整个batch都被mask，不使用mask elif mask_2d.all(dim=1).any(): # 如果某个样本的所有时间步都被mask，至少保留一个时间步不被mask（避免全mask） for i in range(mask_2d.size(0)): if mask_2d[i].all(): mask_2d[i, 0] = False # 至少保留第一个时间步 elif mask is not None and mask.dim() == 2: # 如果已经是2D的，直接使用（假设True表示需要mask） mask_2d = mask.bool() if mask.dtype != torch.bool else mask # 同样检查全mask的情况 if mask_2d.all(): mask_2d = None elif mask_2d.all(dim=1).any(): for i in range(mask_2d.size(0)): if mask_2d[i].all(): mask_2d[i, 0] = False else: mask_2d = None # 编码 encoder_output = self.encoder(encoder_input, mask=mask_2d) # [batch_size, seq_len, hidden_size] if debug and torch.isnan(encoder_output).any(): print(f" ⚠️ encoder_output包含NaN（在encoder后）") # 嵌入已知未来特征 known_future_emb = self.known_future_embedding(known_future_features) # [batch_size, pred_len, hidden_size] if debug and torch.isnan(known_future_emb).any(): print(f" ⚠️ known_future_emb包含NaN（在embedding后）") known_future_emb = self.pos_encoder(known_future_emb) if debug and torch.isnan(known_future_emb).any(): print(f" ⚠️ known_future_emb包含NaN（在pos_encoder后）") # 解码（使用相同的mask_2d） decoder_output = self.decoder( known_future_emb, encoder_output, memory_mask=mask_2d ) # [batch_size, pred_len, hidden_size] if debug and torch.isnan(decoder_output).any(): print(f" ⚠️ decoder_output包含NaN（在decoder后）") # 输出预测 predictions = self.output_layer(decoder_output) # [batch_size, pred_len, num_observed_features] if debug and torch.isnan(predictions).any(): print(f" ⚠️ predictions包含NaN（在output_layer后）") return predictions class PositionalEncoding(nn.Module): """位置编码""" def __init__(self, d_model, dropout=0.1, max_len=5000): super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(p=dropout) position = torch.arange(max_len).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) pe = torch.zeros(max_len, d_model) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) self.register_buffer('pe', pe) def forward(self, x): x = x + self.pe[:, :x.size(1), :] return self.dropout(x)