feat: add mamba and dynamic chunking related code and test code
This commit is contained in:
gameloader
528
models/DC_PatchTST.py
Normal file
528
models/DC_PatchTST.py
Normal file
@ -0,0 +1,528 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from layers.SelfAttention_Family import FullAttention, AttentionLayer
|
||||
|
||||
# 需要 Mamba2 作为外层编码器
|
||||
from mamba_ssm.modules.mamba2 import Mamba2
|
||||
|
||||
|
||||
|
||||
# -------------------- Routing（余弦路由，和论文一致） --------------------
|
||||
class RoutingModule(nn.Module):
|
||||
def __init__(self, d_model):
|
||||
super().__init__()
|
||||
self.q_proj = nn.Linear(d_model, d_model, bias=False)
|
||||
self.k_proj = nn.Linear(d_model, d_model, bias=False)
|
||||
with torch.no_grad():
|
||||
nn.init.eye_(self.q_proj.weight)
|
||||
nn.init.eye_(self.k_proj.weight)
|
||||
self.q_proj.weight._no_reinit = True
|
||||
self.k_proj.weight._no_reinit = True
|
||||
|
||||
def forward(self, x, mask=None):
|
||||
"""
|
||||
x: (B, L, D)
|
||||
mask: (B, L) bool, True=有效
|
||||
返回：
|
||||
boundary_prob: (B, L, 2)
|
||||
boundary_mask: (B, L) bool
|
||||
selected_probs: (B, L, 1)
|
||||
"""
|
||||
B, L, D = x.shape
|
||||
q = F.normalize(self.q_proj(x[:, :-1]), dim=-1) # (B, L-1, D)
|
||||
k = F.normalize(self.k_proj(x[:, 1:]), dim=-1) # (B, L-1, D)
|
||||
cos_sim = (q * k).sum(dim=-1) # (B, L-1)
|
||||
p = torch.clamp((1 - cos_sim) / 2, 0.0, 1.0) # (B, L-1)
|
||||
p = F.pad(p, (1, 0), value=1.0) # 强制首位是边界
|
||||
|
||||
if mask is not None:
|
||||
p = p * mask.float()
|
||||
p[:, 0] = torch.where(mask[:, 0], torch.ones_like(p[:, 0]), p[:, 0])
|
||||
|
||||
boundary_prob = torch.stack([1 - p, p], dim=-1) # (B, L, 2)
|
||||
selected_idx = boundary_prob.argmax(dim=-1)
|
||||
boundary_mask = (selected_idx == 1)
|
||||
if mask is not None:
|
||||
boundary_mask = boundary_mask & mask
|
||||
selected_probs = boundary_prob.gather(-1, selected_idx.unsqueeze(-1)) # (B, L, 1)
|
||||
return boundary_prob, boundary_mask, selected_probs
|
||||
|
||||
|
||||
# -------------------- 选择并右侧零pad（不丢弃、不重复填充） --------------------
|
||||
def select_and_right_pad(x, boundary_mask):
|
||||
"""
|
||||
内存优化版本：减少临时tensor创建
|
||||
x: (B, L, D), boundary_mask: (B, L) bool
|
||||
返回:
|
||||
x_pad: (B, T_max, D)
|
||||
key_padding_mask: (B, T_max) bool, True=有效
|
||||
lengths: (B,)
|
||||
"""
|
||||
B, L, D = x.shape
|
||||
device = x.device
|
||||
lengths = boundary_mask.sum(dim=1) # (B,)
|
||||
T_max = int(lengths.max().item()) if lengths.max() > 0 else 1
|
||||
|
||||
x_pad = x.new_zeros(B, T_max, D)
|
||||
key_padding_mask = torch.zeros(B, T_max, dtype=torch.bool, device=device)
|
||||
|
||||
# 预创建默认索引tensor避免重复创建
|
||||
default_idx = torch.tensor([0], device=device)
|
||||
|
||||
for b in range(B):
|
||||
mask_b = boundary_mask[b]
|
||||
if mask_b.any():
|
||||
idx = mask_b.nonzero(as_tuple=True)[0] # 更高效的nonzero
|
||||
t = idx.numel()
|
||||
x_pad[b, :t] = x[b, idx]
|
||||
key_padding_mask[b, :t] = True
|
||||
else:
|
||||
# 使用预创建的tensor
|
||||
x_pad[b, 0] = x[b, default_idx]
|
||||
key_padding_mask[b, 0] = True
|
||||
|
||||
return x_pad, key_padding_mask, lengths
|
||||
|
||||
|
||||
# -------------------- Mamba2 堆叠（外层编码器） --------------------
|
||||
class Mamba2Encoder(nn.Module):
|
||||
def __init__(self, d_model, depth=4, dropout=0.0):
|
||||
super().__init__()
|
||||
self.layers = nn.ModuleList([Mamba2(d_model=d_model) for _ in range(depth)])
|
||||
self.norm = nn.LayerNorm(d_model)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, x):
|
||||
for layer in self.layers:
|
||||
x = layer(x)
|
||||
x = self.norm(x)
|
||||
x = self.dropout(x)
|
||||
return x
|
||||
|
||||
|
||||
# -------------------- 两层Encoder + DC（变长，不丢信息；以比率约束压缩） --------------------
|
||||
class DCEmbedding2StageVarLen(nn.Module):
|
||||
"""
|
||||
- Stage 0: (B*nvars, L, 1) -> Linear(D0) -> Mamba2(D0) -> Routing -> 选择 -> 扩宽到 D1
|
||||
- Stage 1: (B*nvars, L0_sel, D1) -> Mamba2(D1) -> Routing -> 选择
|
||||
输出：
|
||||
enc_out: (B*nvars, T_max, D1)
|
||||
key_padding_mask: (B*nvars, T_max)
|
||||
n_vars: int
|
||||
aux: dict（含两层ratio loss与边界信息）
|
||||
"""
|
||||
def __init__(self, d_model_out, d_model_stage0, depth_enc0=4, depth_enc1=4, dropout=0.0,
|
||||
target_ratio0=0.25, target_ratio1=0.5):
|
||||
super().__init__()
|
||||
assert d_model_out >= d_model_stage0, "要求 D0 <= D1"
|
||||
self.d0 = d_model_stage0
|
||||
self.d1 = d_model_out
|
||||
|
||||
# 标量 -> D0
|
||||
self.input_proj = nn.Linear(1, self.d0)
|
||||
|
||||
# Stage 0
|
||||
self.enc0 = Mamba2Encoder(self.d0, depth=depth_enc0, dropout=dropout)
|
||||
self.router0 = RoutingModule(self.d0)
|
||||
delta = self.d1 - self.d0
|
||||
self.pad_vec = nn.Parameter(torch.zeros(delta)) if delta > 0 else None
|
||||
self.target_ratio0 = target_ratio0
|
||||
|
||||
# Stage 1
|
||||
self.enc1 = Mamba2Encoder(self.d1, depth=depth_enc1, dropout=dropout)
|
||||
self.router1 = RoutingModule(self.d1)
|
||||
self.target_ratio1 = target_ratio1
|
||||
|
||||
def _expand_width(self, x):
|
||||
if self.pad_vec is None:
|
||||
return x
|
||||
B, L, _ = x.shape
|
||||
return torch.cat([x, self.pad_vec.view(1, 1, -1).expand(B, L, -1)], dim=-1)
|
||||
|
||||
@staticmethod
|
||||
def _ratio_loss(boundary_mask: torch.Tensor, boundary_prob: torch.Tensor, target_ratio: float) -> torch.Tensor:
|
||||
eps = 1e-6
|
||||
F_act = boundary_mask.float().mean(dim=1) # (B,)
|
||||
G_prob = boundary_prob[..., 1].mean(dim=1) # (B,)
|
||||
N = 1.0 / max(target_ratio, eps)
|
||||
loss = N / (N - 1.0 + eps) * (((N - 1.0) * F_act * G_prob) + (1.0 - F_act) * (1.0 - G_prob))
|
||||
return loss.mean()
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: (B, nvars, L)
|
||||
内存优化版本：及时删除中间tensor
|
||||
"""
|
||||
B, nvars, L = x.shape
|
||||
x = x.reshape(B * nvars, L, 1)
|
||||
x = self.input_proj(x) # (B*nvars, L, D0)
|
||||
|
||||
# Stage 0
|
||||
h0 = self.enc0(x) # (B*nvars, L, D0)
|
||||
p0, bm0, _ = self.router0(h0)
|
||||
h0_sel, mask0, len0 = select_and_right_pad(h0, bm0) # (B*nvars, L0_max, D0)
|
||||
|
||||
# 及时删除不需要的tensor
|
||||
del h0
|
||||
|
||||
# h0_sel = self._expand_width(h0_sel) # (B*nvars, L0_max, D1)
|
||||
|
||||
# Stage 1
|
||||
#h1 = self.enc1(h0_sel) # (B*nvars, L0_max, D1)
|
||||
#p1, bm1, _ = self.router1(h1)
|
||||
#bm1 = bm1 & mask0
|
||||
#h1_sel, mask1, len1 = select_and_right_pad(h1, bm1) # (B*nvars, L1_max, D1)
|
||||
|
||||
# 及时删除中间tensor
|
||||
#del h1, h0_sel
|
||||
|
||||
# 计算ratio loss时使用detach避免保存计算图
|
||||
ratio_loss0 = self._ratio_loss(bm0, p0, target_ratio=self.target_ratio0)
|
||||
# ratio_loss1 = self._ratio_loss(bm1, p1, target_ratio=self.target_ratio1)
|
||||
|
||||
# 简化aux字典，只保存必要信息
|
||||
aux = {
|
||||
"stage0": {"boundary_mask": bm0.detach(), "boundary_prob": p0.detach(), "lengths": len0.detach()},
|
||||
# "stage1": {"boundary_mask": bm1.detach(), "boundary_prob": p1.detach(), "lengths": len1.detach()},
|
||||
"ratio_loss0": ratio_loss0,
|
||||
# "ratio_loss1": ratio_loss1,
|
||||
}
|
||||
|
||||
return h0_sel, mask0, nvars, aux
|
||||
|
||||
|
||||
# -------------------- Encoder/EncoderLayer（带 key_padding_mask 透传） --------------------
|
||||
class EncoderLayerWithMask(nn.Module):
|
||||
"""
|
||||
与原EncoderLayer结构一致，但 forward 增加 key_padding_mask，并传入 AttentionLayer。
|
||||
FFN 用简单的 MLP（与常规Transformer一致）。
|
||||
"""
|
||||
def __init__(self, attention: AttentionLayer, d_model, d_ff, dropout=0.1, activation="gelu"):
|
||||
super().__init__()
|
||||
self.attention = attention
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.norm1 = nn.LayerNorm(d_model)
|
||||
self.norm2 = nn.LayerNorm(d_model)
|
||||
|
||||
if activation == "relu":
|
||||
act = nn.ReLU()
|
||||
elif activation == "gelu":
|
||||
act = nn.GELU()
|
||||
else:
|
||||
raise ValueError(f"Unsupported activation: {activation}")
|
||||
|
||||
self.ffn = nn.Sequential(
|
||||
nn.Linear(d_model, d_ff),
|
||||
act,
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(d_ff, d_model),
|
||||
)
|
||||
|
||||
def forward(self, x, attn_mask=None, tau=None, delta=None, key_padding_mask=None):
|
||||
# Multi-head attention with key padding mask
|
||||
attn_out, attn = self.attention(
|
||||
x, x, x, attn_mask, tau=tau, delta=delta, key_padding_mask=key_padding_mask
|
||||
)
|
||||
x = x + self.dropout(attn_out)
|
||||
x = self.norm1(x)
|
||||
|
||||
# FFN
|
||||
y = self.ffn(x)
|
||||
x = x + self.dropout(y)
|
||||
x = self.norm2(x)
|
||||
return x, attn
|
||||
|
||||
|
||||
class EncoderWithMask(nn.Module):
|
||||
"""
|
||||
与原Encoder类似，但 forward 支持 key_padding_mask，并传递给每一层的注意力。
|
||||
"""
|
||||
def __init__(self, attn_layers, norm_layer=None):
|
||||
super().__init__()
|
||||
self.attn_layers = nn.ModuleList(attn_layers)
|
||||
self.norm = norm_layer
|
||||
|
||||
def forward(self, x, attn_mask=None, key_padding_mask=None):
|
||||
attns = []
|
||||
for attn_layer in self.attn_layers:
|
||||
x, attn = attn_layer(x, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
|
||||
attns.append(attn)
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
return x, attns
|
||||
|
||||
|
||||
# -------------------- 门控注意力聚合 + 任务头（不依赖token数；保留信息） --------------------
|
||||
def masked_softmax(logits: torch.Tensor, mask: torch.Tensor, dim: int = -1) -> torch.Tensor:
|
||||
"""
|
||||
logits: (..., T)
|
||||
mask: (..., T) bool, True=有效
|
||||
"""
|
||||
neg_inf = torch.finfo(logits.dtype).min
|
||||
logits = logits.masked_fill(~mask, neg_inf)
|
||||
return torch.softmax(logits, dim=dim)
|
||||
|
||||
class GatedAttnAggregator(nn.Module):
|
||||
"""
|
||||
门控注意力聚合器（可学习查询 + mask softmax + 值端门控）
|
||||
输入: x: (B*, T, D), mask: (B*, T) bool
|
||||
输出: slots: (B*, R, D) 其中 R 为聚合插槽数（可配置）
|
||||
"""
|
||||
def __init__(self, d_model: int, num_slots: int = 4, d_att: int = None, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.d_model = d_model
|
||||
self.R = num_slots
|
||||
self.d_att = d_att or d_model
|
||||
|
||||
# 可学习查询（R个）
|
||||
self.query = nn.Parameter(torch.randn(self.R, self.d_att) / (self.d_att ** 0.5))
|
||||
|
||||
# 线性投影
|
||||
self.key_proj = nn.Linear(d_model, self.d_att)
|
||||
self.val_proj = nn.Linear(d_model, d_model)
|
||||
|
||||
# 值端门控（逐token标量门）
|
||||
self.gate = nn.Sequential(
|
||||
nn.Linear(d_model, d_model // 2),
|
||||
nn.GELU(),
|
||||
nn.Linear(d_model // 2, 1),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, x_bt_t_d: torch.Tensor, mask_bt: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
x_bt_t_d: (B*, T, D)
|
||||
mask_bt: (B*, T) bool
|
||||
return: slots (B*, R, D)
|
||||
"""
|
||||
BStar, T, D = x_bt_t_d.shape
|
||||
K = self.key_proj(x_bt_t_d) # (B*, T, d_att)
|
||||
V = self.val_proj(x_bt_t_d) # (B*, T, D)
|
||||
g = self.gate(x_bt_t_d) # (B*, T, 1)
|
||||
Vg = V * g # 门控后的值
|
||||
|
||||
Q = self.query.unsqueeze(0).expand(BStar, -1, -1) # (B*, R, d_att)
|
||||
|
||||
logits = torch.matmul(Q, K.transpose(1, 2)) / (self.d_att ** 0.5) # (B*, R, T)
|
||||
attn_mask = mask_bt.unsqueeze(1) # (B*, 1, T)
|
||||
attn = masked_softmax(logits, attn_mask, dim=-1)
|
||||
attn = self.dropout(attn)
|
||||
|
||||
slots = torch.matmul(attn, Vg) # (B*, R, D)
|
||||
return slots
|
||||
|
||||
class AttnPoolHeadForecast(nn.Module):
|
||||
"""
|
||||
预测任务头：门控注意力聚合到 R 个slots，再映射到 target_window（pred_len）
|
||||
输出：(B, pred_len, nvars)
|
||||
"""
|
||||
def __init__(self, d_model: int, target_window: int, num_slots: int = 4, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.agg = GatedAttnAggregator(d_model, num_slots=num_slots, dropout=dropout)
|
||||
self.proj = nn.Sequential(
|
||||
nn.LayerNorm(num_slots * d_model),
|
||||
nn.Linear(num_slots * d_model, target_window),
|
||||
)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, enc_out_bt_t_d: torch.Tensor, key_padding_mask_bt: torch.Tensor, n_vars: int, B: int):
|
||||
slots = self.agg(enc_out_bt_t_d, key_padding_mask_bt) # (B*, R, D)
|
||||
slots = slots.reshape(B, n_vars, -1) # (B, nvars, R*D)
|
||||
out = self.proj(self.dropout(slots)) # (B, nvars, pred_len)
|
||||
return out.permute(0, 2, 1) # (B, pred_len, nvars)
|
||||
|
||||
class AttnPoolHeadSeq(nn.Module):
|
||||
"""
|
||||
序列重建头：门控注意力聚合后映射到 seq_len
|
||||
输出：(B, seq_len, nvars)
|
||||
"""
|
||||
def __init__(self, d_model: int, target_window: int, num_slots: int = 4, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.agg = GatedAttnAggregator(d_model, num_slots=num_slots, dropout=dropout)
|
||||
self.proj = nn.Sequential(
|
||||
nn.LayerNorm(num_slots * d_model),
|
||||
nn.Linear(num_slots * d_model, target_window),
|
||||
)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, enc_out_bt_t_d: torch.Tensor, key_padding_mask_bt: torch.Tensor, n_vars: int, B: int):
|
||||
slots = self.agg(enc_out_bt_t_d, key_padding_mask_bt) # (B*, R, D)
|
||||
slots = slots.reshape(B, n_vars, -1) # (B, nvars, R*D)
|
||||
out = self.proj(self.dropout(slots)) # (B, nvars, seq_len)
|
||||
return out.permute(0, 2, 1) # (B, seq_len, nvars)
|
||||
|
||||
class AttnPoolHeadCls(nn.Module):
|
||||
"""
|
||||
分类头：每变量先门控注意力聚合到 R 个slots，拼接所有变量后线性分类。
|
||||
输出：(B, num_class)
|
||||
"""
|
||||
def __init__(self, d_model: int, n_vars: int, num_class: int, num_slots: int = 4, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.agg = GatedAttnAggregator(d_model, num_slots=num_slots, dropout=dropout)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.proj = nn.Sequential(
|
||||
nn.LayerNorm(n_vars * num_slots * d_model),
|
||||
nn.Linear(n_vars * num_slots * d_model, num_class),
|
||||
)
|
||||
self.n_vars = n_vars
|
||||
self.num_slots = num_slots
|
||||
self.d_model = d_model
|
||||
|
||||
def forward(self, enc_out_bt_t_d: torch.Tensor, key_padding_mask_bt: torch.Tensor, n_vars: int, B: int):
|
||||
slots = self.agg(enc_out_bt_t_d, key_padding_mask_bt) # (B*, R, D)
|
||||
slots = slots.reshape(B, n_vars, self.num_slots * self.d_model) # (B, nvars, R*D)
|
||||
flat = self.dropout(slots.reshape(B, -1)) # (B, nvars*R*D)
|
||||
return self.proj(flat)
|
||||
|
||||
|
||||
# -------------------- 主模型：两层DC(比率控制) + 带mask的Encoder + 门控聚合头 --------------------
|
||||
class Transpose(nn.Module):
|
||||
def __init__(self, *dims, contiguous=False):
|
||||
super().__init__()
|
||||
self.dims, self.contiguous = dims, contiguous
|
||||
def forward(self, x):
|
||||
return x.transpose(*self.dims).contiguous() if self.contiguous else x.transpose(*self.dims)
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
PatchTST with DC and masked attention + gated heads:
|
||||
- 用两层 Mamba2 编码器 + 动态分块 替代 PatchEmbedding
|
||||
- DC 使用 ratio loss（target_ratio0/1）控制压缩强度；随层级加深，序列变短，d_model 变大（D0->D1）
|
||||
- 注意力传入 key_padding_mask 屏蔽pad
|
||||
- 头部使用门控注意力聚合（不依赖token数，信息保留更充分）
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, configs,
|
||||
d_model_stage0=None, # D0，默认= d_model // 2
|
||||
depth_enc0=1, depth_enc1=1,
|
||||
target_ratio0=0.25, # 约等于 1/N0
|
||||
target_ratio1=0.5, # 约等于 1/N1
|
||||
agg_slots=4, # 门控聚合的slot数
|
||||
):
|
||||
super().__init__()
|
||||
self.task_name = configs.task_name
|
||||
self.seq_len = configs.seq_len
|
||||
self.pred_len = configs.pred_len
|
||||
self.enc_in = configs.enc_in
|
||||
|
||||
# DC 嵌入
|
||||
D1 = configs.d_model
|
||||
D0 = d_model_stage0 if d_model_stage0 is not None else max(16, D1 // 2)
|
||||
assert D1 >= D0, "要求 D0 <= D1"
|
||||
self.dc_embedding = DCEmbedding2StageVarLen(
|
||||
d_model_out=D1,
|
||||
d_model_stage0=D0,
|
||||
depth_enc0=depth_enc0,
|
||||
depth_enc1=depth_enc1,
|
||||
dropout=configs.dropout,
|
||||
target_ratio0=target_ratio0,
|
||||
target_ratio1=target_ratio1,
|
||||
)
|
||||
|
||||
# 带mask的Encoder
|
||||
attn_layers = [
|
||||
EncoderLayerWithMask(
|
||||
AttentionLayer(
|
||||
FullAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
|
||||
D1, configs.n_heads
|
||||
),
|
||||
d_model=D1,
|
||||
d_ff=configs.d_ff,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation
|
||||
) for _ in range(configs.e_layers)
|
||||
]
|
||||
self.encoder = EncoderWithMask(
|
||||
attn_layers,
|
||||
norm_layer=nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(D1), Transpose(1, 2))
|
||||
)
|
||||
|
||||
# 门控聚合头（与token数无关）
|
||||
if self.task_name in ('long_term_forecast', 'short_term_forecast'):
|
||||
self.head = AttnPoolHeadForecast(D1, self.pred_len, num_slots=agg_slots, dropout=configs.dropout)
|
||||
elif self.task_name in ('imputation', 'anomaly_detection'):
|
||||
self.head = AttnPoolHeadSeq(D1, self.seq_len, num_slots=agg_slots, dropout=configs.dropout)
|
||||
elif self.task_name == 'classification':
|
||||
self.head_cls = AttnPoolHeadCls(D1, n_vars=self.enc_in, num_class=configs.num_class, num_slots=agg_slots, dropout=configs.dropout)
|
||||
|
||||
# --------- 归一化/反归一化 ---------
|
||||
def _pre_norm(self, x):
|
||||
means = x.mean(1, keepdim=True).detach()
|
||||
x = x - means
|
||||
stdev = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False) + 1e-5)
|
||||
x = x / stdev
|
||||
return x, means, stdev
|
||||
|
||||
def _denorm(self, y, means, stdev, length):
|
||||
return y * (stdev[:, 0, :].unsqueeze(1).repeat(1, length, 1)) + \
|
||||
(means[:, 0, :].unsqueeze(1).repeat(1, length, 1))
|
||||
|
||||
# --------- DC + Transformer Encoder（携带 key_padding_mask） ----------
|
||||
def _embed_and_encode(self, x_enc):
|
||||
"""
|
||||
x_enc: (B, L, C)
|
||||
返回：
|
||||
enc_out: (B*nvars, T_max, D1)
|
||||
n_vars: int
|
||||
key_padding_mask: (B*nvars, T_max)
|
||||
aux: dict
|
||||
"""
|
||||
B, L, C = x_enc.shape
|
||||
x_vars = x_enc.permute(0, 2, 1) # (B, nvars, L)
|
||||
enc_out, key_padding_mask, n_vars, aux = self.dc_embedding(x_vars)
|
||||
enc_out, _ = self.encoder(enc_out, attn_mask=None, key_padding_mask=key_padding_mask)
|
||||
return enc_out, n_vars, key_padding_mask, B, aux
|
||||
|
||||
# --------- 各任务前向 ---------
|
||||
def forecast(self, x_enc, x_mark_enc, x_dec, x_mark_dec):
|
||||
x_enc, means, stdev = self._pre_norm(x_enc)
|
||||
enc_out, n_vars, key_padding_mask, B, aux = self._embed_and_encode(x_enc)
|
||||
dec_out = self.head(enc_out, key_padding_mask, n_vars, B) # (B, pred_len, nvars)
|
||||
dec_out = self._denorm(dec_out, means, stdev, self.pred_len)
|
||||
return dec_out, aux
|
||||
|
||||
def imputation(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask):
|
||||
means = torch.sum(x_enc, dim=1) / torch.sum(mask == 1, dim=1)
|
||||
means = means.unsqueeze(1).detach()
|
||||
x = x_enc - means
|
||||
x = x.masked_fill(mask == 0, 0)
|
||||
stdev = torch.sqrt(torch.sum(x * x, dim=1) / torch.sum(mask == 1, dim=1) + 1e-5)
|
||||
stdev = stdev.unsqueeze(1).detach()
|
||||
x = x / stdev
|
||||
|
||||
enc_out, n_vars, key_padding_mask, B, aux = self._embed_and_encode(x)
|
||||
dec_out = self.head(enc_out, key_padding_mask, n_vars, B) # (B, seq_len, nvars)
|
||||
dec_out = self._denorm(dec_out, means, stdev, self.seq_len)
|
||||
return dec_out, aux
|
||||
|
||||
def anomaly_detection(self, x_enc):
|
||||
x_enc, means, stdev = self._pre_norm(x_enc)
|
||||
enc_out, n_vars, key_padding_mask, B, aux = self._embed_and_encode(x_enc)
|
||||
dec_out = self.head(enc_out, key_padding_mask, n_vars, B) # (B, seq_len, nvars)
|
||||
dec_out = self._denorm(dec_out, means, stdev, self.seq_len)
|
||||
return dec_out, aux
|
||||
|
||||
def classification(self, x_enc, x_mark_enc):
|
||||
x_enc, _, _ = self._pre_norm(x_enc)
|
||||
enc_out, n_vars, key_padding_mask, B, aux = self._embed_and_encode(x_enc)
|
||||
logits = self.head_cls(enc_out, key_padding_mask, n_vars, B) # (B, num_class)
|
||||
return logits, aux
|
||||
|
||||
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):
|
||||
if self.task_name in ('long_term_forecast', 'short_term_forecast'):
|
||||
dec_out, aux = self.forecast(x_enc, x_mark_enc, x_dec, x_mark_dec)
|
||||
return dec_out[:, -self.pred_len:, :], aux # [B, L, D], aux含ratio losses
|
||||
if self.task_name == 'imputation':
|
||||
dec_out, aux = self.imputation(x_enc, x_mark_enc, x_dec, x_mark_dec, mask)
|
||||
return dec_out, aux
|
||||
if self.task_name == 'anomaly_detection':
|
||||
dec_out, aux = self.anomaly_detection(x_enc)
|
||||
return dec_out, aux
|
||||
if self.task_name == 'classification':
|
||||
logits, aux = self.classification(x_enc, x_mark_enc)
|
||||
return logits, aux
|
||||
return None, None
|
||||
Reference in New Issue
Block a user