File size: 5,536 Bytes
ddd3436 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | from typing import List, Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from model.TS_encoder.ts_encoder import TimeSeriesEncoder
class TS_Model(nn.Module):
"""Model for time series pretraining with masked reconstruction and anomaly detection."""
def __init__(self, config_t, **kwargs):
super().__init__()
self.config = config_t
ts_config = self.config
self.ts_encoder = TimeSeriesEncoder(
d_model=ts_config.d_model,
d_proj=ts_config.d_proj,
patch_size=ts_config.patch_size,
num_layers=ts_config.num_layers,
num_heads=ts_config.num_heads,
d_ff_dropout=ts_config.d_ff_dropout,
use_rope=ts_config.use_rope,
num_features=ts_config.num_features,
activation=ts_config.activation
)
self.d_proj=ts_config.d_proj
self.patch_size=ts_config.patch_size
self.token_hidden_size=512
self.MAX_L=5000
self.reconstruction_head = nn.Sequential(
nn.Linear(self.d_proj, self.d_proj* 4),
nn.GELU(),
nn.Dropout(0.1),
nn.Linear(self.d_proj* 4, self.d_proj * 4),
nn.GELU(),
nn.Dropout(0.1),
nn.Linear(self.d_proj * 4, 1)
)
# Anomaly detection head
self.anomaly_head = nn.Sequential(
nn.Linear(self.d_proj, self.d_proj // 2),
nn.GELU(),
nn.Dropout(0.1),
nn.Linear(self.d_proj // 2, 2) # binary classification
)
def forward(self,time_series: torch.Tensor, mask: Optional[torch.Tensor] = None):
patch_embeddings,local_embeddings,full_mask=self.ts_encoder(time_series,mask)
return patch_embeddings,local_embeddings,full_mask
def masked_reconstruction_loss(self,
local_embeddings: torch.Tensor,
original_time_series: torch.Tensor,
mask: torch.Tensor) -> torch.Tensor:
"""Compute masked reconstruction loss."""
batch_size, seq_len,num_features = original_time_series.shape
patch_size = self.patch_size
mask = mask.bool()
reconstructed = self.reconstruction_head(local_embeddings).squeeze(-1) # [B, seq_len, 1]
mask_expanded = mask.unsqueeze(-1).expand(-1, -1, num_features)
reconstruction_loss = F.mse_loss(
reconstructed[mask_expanded],
original_time_series[mask_expanded],
)
error = (reconstructed - original_time_series).abs()
return reconstruction_loss,error
def weighted_reconstruction_loss(
self,
local_embeddings: torch.Tensor,
original_time_series: torch.Tensor,
mask: torch.Tensor,
labels: torch.Tensor,
):
"""
Compute weighted reconstruction loss with optional mask weakening,
and return the self-similarity matrix of reconstructed features.
Args:
local_embeddings: [B, seq_len, d_proj]
original_time_series: [B, seq_len, 1] or [B, seq_len, C]
mask: [B, seq_len], bool or float (1 = masked)
mask_weight: weight for masked positions (e.g., 0.5 to weaken)
Returns:
total_loss: scalar loss
error: absolute error per position [B, seq_len, C]
sim_matrix: self-similarity matrix of reconstructed embeddings [B, seq_len, seq_len]
"""
batch_size, seq_len, num_f = original_time_series.shape
device = original_time_series.device
# Ensure mask is boolean
if not mask.dtype == torch.bool:
mask = mask > 0.5 # threshold if float mask
# Reconstruct time series from embeddings
reconstructed = self.reconstruction_head(local_embeddings).squeeze(-1) # [B, seq_len, C]
effective_mask = mask.clone() # [B, L]
if labels is not None:
labels = labels.bool() # [B, L]
effective_mask = effective_mask & (~labels)
flat_mask = effective_mask.view(-1)
reconstruction_loss = F.mse_loss(
reconstructed.reshape(-1, num_f)[flat_mask],
original_time_series.reshape(-1, num_f)[flat_mask],
)
emb_norm = F.normalize(local_embeddings.flatten(2), p=2, dim=-1) # [B, seq_len, d_proj]
# sim_matrix = torch.bmm(emb_norm, emb_norm.transpose(1, 2)) # [B, seq_len, seq_len]
return reconstruction_loss, reconstructed
def anomaly_detection_loss(self,
local_embeddings: torch.Tensor,
labels: torch.Tensor) -> torch.Tensor:
"""Compute anomaly detection loss for each timestep."""
# Project local embeddings to anomaly scores
logits = self.anomaly_head(local_embeddings) # [B, seq_len, 2]
logits = torch.mean(logits, dim=-2)
# Reshape for loss computation
attention_mask= (labels != -1)
# Create mask for valid labels (not padding)
# Compute loss only on valid timesteps
if attention_mask.sum() > 0:
anomaly_loss = F.cross_entropy(
logits[attention_mask],
labels[attention_mask].long()
)
else:
anomaly_loss = torch.tensor(0.0, device=logits.device)
return anomaly_loss,logits |