FlowVortex/SymTime
113
1from typing import Optional, Union, Tuple, Callable2import math3 4import numpy as np5import torch6from torch import nn7from torch import Tensor8import torch.nn.functional as F9from einops import rearrange10 11 12def get_activation_fn(activation: Union[str, Callable]) -> nn.Module:13 """14 Select the activation function to use.15 16 Parameters17 ----------18 activation : Union[str, Callable]19 The activation specification to resolve. It can be a string such as20 "relu" or "gelu", or a callable that returns an activation module.21 22 Return23 ------24 nn.Module25 The corresponding activation module instance.26 """27 if callable(activation):28 return activation()29 elif activation.lower() == "relu":30 return nn.ReLU()31 elif activation.lower() == "gelu":32 return nn.GELU()33 raise ValueError(34 f'{activation} is not available. You can use "relu", "gelu", or a callable'35 )36 37 38class Transpose(nn.Module):39 """Transpose the dimensions of the input tensor.40 41 Parameters42 ----------43 *dims : int44 The dimensions passed to `torch.Tensor.transpose`.45 contiguous : bool, optional46 Whether to return a contiguous tensor after transposing, by default False.47 48 Return49 ------50 Tensor51 The transposed tensor.52 """53 54 def __init__(self, *dims, contiguous=False) -> None:55 super().__init__()56 self.dims, self.contiguous = dims, contiguous57 58 def forward(self, x: Tensor) -> Tensor:59 if self.contiguous:60 return x.transpose(*self.dims).contiguous()61 else:62 return x.transpose(*self.dims)63 64 65class PositionalEmbedding(nn.Module):66 """Adding the positional encoding to the input for Transformer"""67 68 def __init__(self, hidden_size: int, max_len: int = 5000) -> None:69 super(PositionalEmbedding, self).__init__()70 71 # Calculate the positional encoding once in the logarithmic space.72 pe = torch.zeros(73 max_len, hidden_size74 ).float() # Initialize a tensor of zeros with shape (max_len, hidden_size) to store positional encodings75 pe.requires_grad = (76 False # Positional encodings do not require gradients as they are fixed77 )78 79 position = (80 torch.arange(0, max_len).float().unsqueeze(1)81 ) # Generate a sequence from 0 to max_len-1 and add a dimension at the 1st axis82 div_term = (83 torch.arange(0, hidden_size, 2).float() * -(math.log(10000.0) / hidden_size)84 ).exp() # Calculate the divisor term in the positional encoding formula85 86 pe[:, 0::2] = torch.sin(87 position * div_term88 ) # Apply the sine function to the even columns of the positional encoding matrix89 pe[:, 1::2] = torch.cos(90 position * div_term91 ) # Apply the cosine function to the odd columns of the positional encoding matrix92 93 pe = pe.unsqueeze(94 095 ) # Add a batch dimension, changing the shape to (1, max_len, hidden_size)96 self.register_buffer(97 "pe", pe98 ) # Register the positional encodings as a buffer, which will not be updated as model parameters99 100 def forward(self, x: Tensor) -> Tensor:101 # Return the first max_len positional encodings that match the length of input x102 return x + self.pe[:, : x.size(1)]103 104 105class TSTEncoder(nn.Module):106 """Time series encoder backbone of SymTime"""107 108 def __init__(109 self,110 patch_size: int = 16,111 num_layers: int = 3,112 hidden_size: int = 128,113 num_heads: int = 16,114 d_k: int = None,115 d_v: int = None,116 d_ff: int = 256,117 norm: str = "BatchNorm",118 attn_dropout: float = 0.0,119 dropout: float = 0.0,120 act: str = "gelu",121 store_attn: bool = False,122 pre_norm: bool = False,123 ) -> None:124 super().__init__()125 # The Linear layer to project the input patches to the model dimension126 self.W_P = nn.Linear(patch_size, hidden_size)127 128 # Positional encoding129 self.pe = PositionalEmbedding(hidden_size=hidden_size)130 131 # Residual dropout132 self.dropout = nn.Dropout(dropout)133 134 # Create the [CLS] token135 self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size))136 self.cls_mask = nn.Parameter(torch.ones(1, 1).bool(), requires_grad=False)137 138 # Create the encoder layer of the model backbone139 self.layers = nn.ModuleList(140 [141 TSTEncoderLayer(142 hidden_size=hidden_size,143 num_heads=num_heads,144 d_k=d_k,145 d_v=d_v,146 d_ff=d_ff,147 norm=norm,148 attn_dropout=attn_dropout,149 dropout=dropout,150 activation=act,151 pre_norm=pre_norm,152 store_attn=store_attn,153 )154 for _ in range(num_layers)155 ]156 )157 158 # model params init159 self.apply(self._init_weights)160 161 def _init_weights(self, m: nn.Module) -> None:162 """model params init through apply methods"""163 if isinstance(m, nn.Linear):164 nn.init.xavier_uniform_(m.weight)165 if isinstance(m, nn.Linear) and m.bias is not None:166 nn.init.constant_(m.bias, 0)167 elif isinstance(m, nn.LayerNorm):168 nn.init.constant_(m.bias, 0)169 nn.init.constant_(m.weight, 1.0)170 171 def forward(172 self,173 x: Tensor, # x: [batch_size, patch_num, patch_size]174 attn_mask: Optional[Tensor] = None, # attn_mask: [batch, num_patch]175 return_cls_token: bool = True, # whether to return the CLS token176 ) -> Tensor:177 """ """178 batch_size = x.size(0)179 180 # Input patching embedding181 x = self.W_P(x) # x: [batch_size, patch_num, model_dim]182 183 # Add the [CLS] token184 cls_token = self.cls_token.expand(batch_size, -1, -1)185 x = torch.cat([cls_token, x], dim=1)186 # adjust the attn mask187 if attn_mask is not None:188 attn_mask = torch.cat(189 [self.cls_mask.expand(batch_size, -1), attn_mask], dim=1190 )191 192 # Add the positional embedding193 x = self.pe(x)194 x = self.dropout(x) # x: [batch_size, patch_num, hidden_size]195 196 for mod in self.layers:197 x = mod(x, attn_mask=attn_mask)198 199 if not return_cls_token:200 # If not returning the CLS token, remove it from the output201 return x[:, 1:, :]202 203 return x204 205 206class TSTEncoderLayer(nn.Module):207 """Patch-based Transformer module sublayer"""208 209 def __init__(210 self,211 hidden_size: int,212 num_heads: int,213 d_k: int = None,214 d_v: int = None,215 d_ff: int = 256,216 store_attn: int = False,217 norm: str = "BatchNorm",218 attn_dropout: float = 0.0,219 dropout: float = 0.0,220 bias: bool = True,221 activation: str = "gelu",222 pre_norm: bool = False,223 ) -> None:224 super(TSTEncoderLayer, self).__init__()225 226 assert (227 not hidden_size % num_heads228 ), f"hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})"229 # If not specified, the number of heads is divided230 d_k = hidden_size // num_heads if d_k is None else d_k231 d_v = hidden_size // num_heads if d_v is None else d_v232 233 # Create the multi-head attention234 self.self_attn = MultiHeadAttention(235 hidden_size,236 num_heads,237 d_k,238 d_v,239 attn_dropout=attn_dropout,240 proj_dropout=dropout,241 )242 243 # Add & Norm244 self.dropout_attn = nn.Dropout(dropout)245 if "batch" in norm.lower():246 self.norm_attn = nn.Sequential(247 Transpose(1, 2), nn.BatchNorm1d(hidden_size), Transpose(1, 2)248 )249 else:250 self.norm_attn = nn.LayerNorm(hidden_size)251 252 # Position-wise Feed-Forward253 self.ff = nn.Sequential(254 nn.Linear(hidden_size, d_ff, bias=bias),255 get_activation_fn(activation),256 nn.Dropout(dropout),257 nn.Linear(d_ff, hidden_size, bias=bias),258 )259 260 # Add & Norm261 self.dropout_ffn = nn.Dropout(dropout)262 if "batch" in norm.lower():263 self.norm_ffn = nn.Sequential(264 Transpose(1, 2), nn.BatchNorm1d(hidden_size), Transpose(1, 2)265 )266 else:267 self.norm_ffn = nn.LayerNorm(hidden_size)268 269 # use pre-norm or not270 self.pre_norm = pre_norm271 self.store_attn = store_attn272 self.attn = None273 274 def forward(275 self, src: Tensor, attn_mask: Optional[Tensor] = None276 ) -> Union[Tuple[Tensor, Tensor], Tensor]:277 """Multi-Head attention sublayer"""278 279 # Whether to use pre-norm for attention layer280 if self.pre_norm:281 src = self.norm_attn(src)282 283 # Multi-Head attention284 src2, attn = self.self_attn(src, src, src, attn_mask=attn_mask)285 if self.store_attn:286 self.attn = attn287 288 # Add: residual connection with residual dropout289 src = src + self.dropout_attn(src2)290 if not self.pre_norm:291 src = self.norm_attn(src)292 293 # Whether to use pre-norm for ffn layer294 if self.pre_norm:295 src = self.norm_ffn(src)296 297 # Position-wise Feed-Forward298 src2 = self.ff(src)299 300 # Add: residual connection with residual dropout301 src = src + self.dropout_ffn(src2)302 if not self.pre_norm:303 src = self.norm_ffn(src)304 305 return src306 307 308class MultiHeadAttention(nn.Module):309 """Multi-head attention mechanism layer"""310 311 def __init__(312 self,313 hidden_size: int,314 num_heads: int,315 d_k: int = None,316 d_v: int = None,317 attn_dropout: float = 0.0,318 proj_dropout: float = 0.0,319 qkv_bias: bool = True,320 ) -> None:321 """Multi Head Attention Layer322 Input shape:323 Q: [batch_size (bs) x max_q_len x hidden_size]324 K, V: [batch_size (bs) x q_len x hidden_size]325 mask: [q_len x q_len]326 """327 super().__init__()328 d_k = hidden_size // num_heads if d_k is None else d_k329 d_v = hidden_size // num_heads if d_v is None else d_v330 331 self.num_heads, self.d_k, self.d_v = num_heads, d_k, d_v332 333 self.W_Q = nn.Linear(hidden_size, d_k * num_heads, bias=qkv_bias)334 self.W_K = nn.Linear(hidden_size, d_k * num_heads, bias=qkv_bias)335 self.W_V = nn.Linear(hidden_size, d_v * num_heads, bias=qkv_bias)336 337 # Scaled Dot-Product Attention (multiple heads)338 self.sdp_attn = _ScaledDotProductAttention(339 hidden_size, num_heads, attn_dropout=attn_dropout340 )341 342 # Project output343 self.to_out = nn.Sequential(344 nn.Linear(num_heads * d_v, hidden_size), nn.Dropout(proj_dropout)345 )346 347 def forward(348 self,349 q: Tensor,350 k: Optional[Tensor] = None,351 v: Optional[Tensor] = None,352 attn_mask: Optional[Tensor] = None,353 ):354 bs = q.size(0)355 if k is None:356 k = q357 if v is None:358 v = q359 360 # Linear (+ split in multiple heads)361 q_s = self.W_Q(q).view(bs, -1, self.num_heads, self.d_k).transpose(1, 2)362 k_s = self.W_K(k).view(bs, -1, self.num_heads, self.d_k).permute(0, 2, 3, 1)363 v_s = self.W_V(v).view(bs, -1, self.num_heads, self.d_v).transpose(1, 2)364 365 # Apply Scaled Dot-Product Attention (multiple heads)366 output, attn_weights = self.sdp_attn(q_s, k_s, v_s, attn_mask=attn_mask)367 368 # back to the original inputs dimensions369 output = (370 output.transpose(1, 2).contiguous().view(bs, -1, self.num_heads * self.d_v)371 )372 output = self.to_out(output)373 374 return output, attn_weights375 376 377class _ScaledDotProductAttention(nn.Module):378 r"""Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual attention from previous layer379 (Realformer: Transformer likes residual attention by He et al, 2020) and locality self sttention (Vision Transformer for Small-Size Datasets380 by Lee et al, 2021)"""381 382 def __init__(383 self,384 hidden_size: int,385 num_heads: int,386 attn_dropout: float = 0.0,387 res_attention: bool = False,388 ):389 super().__init__()390 self.attn_dropout = nn.Dropout(attn_dropout)391 self.res_attention = res_attention392 head_dim = hidden_size // num_heads393 self.scale = nn.Parameter(torch.tensor(head_dim**-0.5), requires_grad=False)394 395 def forward(396 self, q: Tensor, k: Tensor, v: Tensor, attn_mask: Optional[Tensor] = None397 ) -> Union[Tuple[Tensor, Tensor, Tensor], Tuple[Tensor, Tensor]]:398 """399 :param q: [batch_size, num_heads, num_token, d_k]400 :param k: [batch_size, num_heads, d_k, num_token]401 :param v: [batch_size, num_heads, num_token, d_k]402 :param attn_mask: [batch_size, num_heads, num_token]403 """404 405 # Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence406 attn_scores = torch.matmul(q, k) * self.scale407 408 # Attention mask (optional)409 if (410 attn_mask is not None411 ): # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len412 attn_mask = rearrange(attn_mask, "b i -> b 1 i 1") * rearrange(413 attn_mask, "b i -> b 1 1 i"414 )415 if attn_mask.dtype == torch.bool:416 attn_scores.masked_fill_(attn_mask, -np.inf)417 else:418 attn_scores += attn_mask419 420 # normalize the attention weights421 attn_weights = F.softmax(attn_scores, dim=-1)422 attn_weights = self.attn_dropout(attn_weights)423 424 # compute the new values given the attention weights425 output = torch.matmul(attn_weights, v)426 427 return output, attn_weights428