| | """ |
| | 1D Vision Transformer for time-series signals. |
| | |
| | Patchify modes: |
| | - lead_wise=0: 1D patchify (all channels in one patch), no lead embedding |
| | - lead_wise=1: 2D patchify (channel groups), with lead embedding by default |
| | """ |
| |
|
| | import torch |
| | import torch.nn as nn |
| | from einops import rearrange |
| |
|
| |
|
| | class DropPath(nn.Module): |
| | def __init__(self, drop_prob: float, scale_by_keep: bool = True): |
| | super().__init__() |
| | self.drop_prob = drop_prob |
| | self.scale_by_keep = scale_by_keep |
| |
|
| | def forward(self, x): |
| | if self.drop_prob <= 0. or not self.training: |
| | return x |
| | keep_prob = 1 - self.drop_prob |
| | shape = (x.shape[0],) + (1,) * (x.ndim - 1) |
| | random_tensor = x.new_empty(shape).bernoulli_(keep_prob) |
| | if keep_prob > 0.0 and self.scale_by_keep: |
| | random_tensor.div_(keep_prob) |
| | return x * random_tensor |
| |
|
| |
|
| | class PreNorm(nn.Module): |
| | def __init__(self, dim: int, fn: nn.Module): |
| | super().__init__() |
| | self.norm = nn.LayerNorm(dim) |
| | self.fn = fn |
| |
|
| | def forward(self, x, **kwargs): |
| | return self.fn(self.norm(x), **kwargs) |
| |
|
| |
|
| | class FeedForward(nn.Module): |
| | def __init__(self, input_dim: int, output_dim: int, hidden_dim: int, drop_out_rate=0.): |
| | super().__init__() |
| | self.net = nn.Sequential( |
| | nn.Linear(input_dim, hidden_dim), |
| | nn.GELU(), |
| | nn.Dropout(drop_out_rate), |
| | nn.Linear(hidden_dim, output_dim), |
| | nn.Dropout(drop_out_rate) |
| | ) |
| |
|
| | def forward(self, x): |
| | return self.net(x) |
| |
|
| |
|
| | class Attention(nn.Module): |
| | def __init__(self, input_dim: int, output_dim: int, heads: int = 8, dim_head: int = 64, |
| | qkv_bias: bool = True, drop_out_rate: float = 0., attn_drop_out_rate: float = 0.): |
| | super().__init__() |
| | inner_dim = dim_head * heads |
| | project_out = not (heads == 1 and dim_head == input_dim) |
| |
|
| | self.heads = heads |
| | self.scale = dim_head ** -0.5 |
| | self.attend = nn.Softmax(dim=-1) |
| | self.dropout = nn.Dropout(attn_drop_out_rate) |
| | self.to_qkv = nn.Linear(input_dim, inner_dim * 3, bias=qkv_bias) |
| |
|
| | if project_out: |
| | self.to_out = nn.Sequential(nn.Linear(inner_dim, output_dim), nn.Dropout(drop_out_rate)) |
| | else: |
| | self.to_out = nn.Identity() |
| |
|
| | def forward(self, x): |
| | qkv = self.to_qkv(x).chunk(3, dim=-1) |
| | q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv) |
| | dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale |
| | attn = self.attend(dots) |
| | attn = self.dropout(attn) |
| | out = torch.matmul(attn, v) |
| | out = rearrange(out, 'b h n d -> b n (h d)') |
| | return self.to_out(out) |
| |
|
| |
|
| | class TransformerBlock(nn.Module): |
| | def __init__(self, input_dim: int, output_dim: int, hidden_dim: int, heads: int = 8, |
| | dim_head: int = 32, qkv_bias: bool = True, drop_out_rate: float = 0., |
| | attn_drop_out_rate: float = 0., drop_path_rate: float = 0.): |
| | super().__init__() |
| | attn = Attention(input_dim, output_dim, heads, dim_head, qkv_bias, drop_out_rate, attn_drop_out_rate) |
| | self.attn = PreNorm(input_dim, attn) |
| | self.droppath1 = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity() |
| |
|
| | ff = FeedForward(output_dim, output_dim, hidden_dim, drop_out_rate) |
| | self.ff = PreNorm(output_dim, ff) |
| | self.droppath2 = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity() |
| |
|
| | def forward(self, x): |
| | x = self.droppath1(self.attn(x)) + x |
| | x = self.droppath2(self.ff(x)) + x |
| | return x |
| |
|
| |
|
| | class ViT(nn.Module): |
| | def __init__(self, |
| | num_leads: int, |
| | seq_len: int, |
| | patch_size: int, |
| | lead_wise=0, |
| | patch_size_ch=4, |
| | use_lead_embedding: bool = True, |
| | width: int = 768, |
| | depth: int = 12, |
| | mlp_dim: int = 3072, |
| | heads: int = 12, |
| | dim_head: int = 64, |
| | qkv_bias: bool = True, |
| | drop_out_rate: float = 0., |
| | attn_drop_out_rate: float = 0., |
| | drop_path_rate: float = 0., |
| | **kwargs): |
| | super().__init__() |
| | assert seq_len % patch_size == 0 |
| | num_patches = seq_len // patch_size |
| | self.lead_wise = lead_wise |
| | self.use_lead_embedding = use_lead_embedding |
| |
|
| | if lead_wise == 0: |
| | self.to_patch_embedding = nn.Conv1d(num_leads, width, kernel_size=patch_size, stride=patch_size, bias=False) |
| | self.pos_embedding = nn.Parameter(torch.randn(1, num_patches, width)) |
| | else: |
| | self.to_patch_embedding = nn.Conv2d(1, width, kernel_size=(patch_size_ch, patch_size), |
| | stride=(patch_size_ch, patch_size), bias=False) |
| | self.pos_embedding = nn.Parameter(torch.randn(1, num_patches * num_leads // patch_size_ch, width)) |
| | if use_lead_embedding: |
| | self.lead_emb = nn.Embedding(num_leads // patch_size_ch, width) |
| | else: |
| | self.lead_emb = None |
| |
|
| | self.dropout = nn.Dropout(drop_out_rate) |
| | self.depth = depth |
| | self.width = width |
| |
|
| | drop_path_rate_list = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] |
| | for i in range(depth): |
| | block = TransformerBlock(width, width, mlp_dim, heads, dim_head, qkv_bias, |
| | drop_out_rate, attn_drop_out_rate, drop_path_rate_list[i]) |
| | self.add_module(f'block{i}', block) |
| |
|
| | self.norm = nn.LayerNorm(width) |
| | self.head = nn.Identity() |
| |
|
| | def _patchify_and_embed(self, series: torch.Tensor) -> torch.Tensor: |
| | """Patchify input and add positional/lead embeddings. [B,C,T] -> [B,N,D]""" |
| | if self.lead_wise == 0: |
| | x = self.to_patch_embedding(series) |
| | x = rearrange(x, 'b c n -> b n c') |
| | x = x + self.pos_embedding[:, :x.size(1), :].to(x.device) |
| | else: |
| | x = self.to_patch_embedding(series.unsqueeze(1)) |
| | Lr, Nt = x.shape[-2], x.shape[-1] |
| | x = rearrange(x, 'b c lr nt -> b (lr nt) c') |
| | x = x + self.pos_embedding[:, :x.size(1), :].to(x.device) |
| | if self.use_lead_embedding and self.lead_emb is not None: |
| | row_ids = torch.arange(Lr, device=x.device).repeat_interleave(Nt) |
| | x = x + self.lead_emb(row_ids)[None, :, :] |
| | return x |
| |
|
| | def forward_encoding(self, series: torch.Tensor) -> torch.Tensor: |
| | """Encode series. Returns [B,D] (mean pooled).""" |
| | x = self._patchify_and_embed(series) |
| | x = self.dropout(x) |
| | for i in range(self.depth): |
| | x = getattr(self, f'block{i}')(x) |
| | x = x.mean(dim=1) |
| | return self.norm(x) |
| |
|
| | def forward(self, series): |
| | x = self.forward_encoding(series) |
| | return self.head(x) |
| |
|
| | def reset_head(self, num_classes=1): |
| | del self.head |
| | self.head = nn.Linear(self.width, num_classes) |
| |
|
| |
|
| | def vit_nano(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs): |
| | return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size, |
| | width=128, depth=6, heads=4, mlp_dim=512, **kwargs) |
| |
|
| |
|
| | def vit_tiny(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs): |
| | return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size, |
| | width=192, depth=12, heads=3, mlp_dim=768, **kwargs) |
| |
|
| |
|
| | def vit_small(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs): |
| | return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size, |
| | width=384, depth=12, heads=6, mlp_dim=1536, **kwargs) |
| |
|
| |
|
| | def vit_middle(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs): |
| | return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size, |
| | width=512, depth=12, heads=8, mlp_dim=2048, **kwargs) |
| |
|
| |
|
| | def vit_base(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs): |
| | return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size, |
| | width=768, depth=12, heads=12, mlp_dim=3072, **kwargs) |
| |
|
