enver1323/patchtst-classification-ecg

1from typing import Optional2 3import torch4from torch import nn, Tensor5import torch.nn.functional as F6 7from huggingface_hub import PyTorchModelHubMixin8 9 10class Transpose(nn.Module):11    def __init__(self, *dims, contiguous=False):12        super(Transpose, self).__init__()13        self.dims, self.contiguous = dims, contiguous14 15    def forward(self, x):16        if self.contiguous:17            return x.transpose(*self.dims).contiguous()18        else:19            return x.transpose(*self.dims)20 21    def __repr__(self):22        if self.contiguous:23            return f"{self.__class__.__name__}(dims={', '.join([str(d) for d in self.dims])}).contiguous()"24        else:25            return (26                f"{self.__class__.__name__}({', '.join([str(d) for d in self.dims])})"27            )28 29 30pytorch_acts = [31    nn.ELU,32    nn.LeakyReLU,33    nn.PReLU,34    nn.ReLU,35    nn.ReLU6,36    nn.SELU,37    nn.CELU,38    nn.GELU,39    nn.Sigmoid,40    nn.Softplus,41    nn.Tanh,42    nn.Softmax,43]44pytorch_act_names = [a.__name__.lower() for a in pytorch_acts]45 46 47def get_act_fn(act, **act_kwargs):48    if act is None:49        return50    elif isinstance(act, nn.Module):51        return act52    elif callable(act):53        return act(**act_kwargs)54    idx = pytorch_act_names.index(act.lower())55    return pytorch_acts[idx](**act_kwargs)56 57 58class RevIN(nn.Module):59    def __init__(60        self,61        c_in: int,62        affine: bool = True,63        subtract_last: bool = False,64        dim: int = 2,65        eps: float = 1e-5,66    ):67        super().__init__()68        self.c_in, self.affine, self.subtract_last, self.dim, self.eps = (69            c_in,70            affine,71            subtract_last,72            dim,73            eps,74        )75        if self.affine:76            self.weight = nn.Parameter(torch.ones(1, c_in, 1))77            self.bias = nn.Parameter(torch.zeros(1, c_in, 1))78 79    def forward(self, x: Tensor, mode: Tensor):80        if mode:81            return self.normalize(x)82        else:83            return self.denormalize(x)84 85    def normalize(self, x):86        if self.subtract_last:87            self.sub = x[..., -1].unsqueeze(-1).detach()88        else:89            self.sub = torch.mean(x, dim=-1, keepdim=True).detach()90        self.std = (91            torch.std(x, dim=-1, keepdim=True, unbiased=False).detach() + self.eps92        )93        if self.affine:94            x = x.sub(self.sub)95            x = x.div(self.std)96            x = x.mul(self.weight)97            x = x.add(self.bias)98            return x99        else:100            x = x.sub(self.sub)101            x = x.div(self.std)102            return x103 104    def denormalize(self, x):105        if self.affine:106            x = x.sub(self.bias)107            x = x.div(self.weight)108            x = x.mul(self.std)109            x = x.add(self.sub)110            return x111        else:112            x = x.mul(self.std)113            x = x.add(self.sub)114            return x115 116 117class MovingAverage(nn.Module):118    def __init__(119        self,120        kernel_size: int,121    ):122        super().__init__()123        padding_left = (kernel_size - 1) // 2124        padding_right = kernel_size - padding_left - 1125        self.padding = torch.nn.ReplicationPad1d((padding_left, padding_right))126        self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=1)127 128    def forward(self, x: Tensor):129        return self.avg(self.padding(x))130 131 132class SeriesDecomposition(nn.Module):133    def __init__(134        self,135        kernel_size: int,  # the size of the window136    ):137        super().__init__()138        self.moving_avg = MovingAverage(kernel_size)139 140    def forward(self, x: Tensor):141        moving_mean = self.moving_avg(x)142        residual = x - moving_mean143        return residual, moving_mean144 145 146class _ScaledDotProductAttention(nn.Module):147    def __init__(self, d_model, n_heads, attn_dropout=0.0, res_attention=False):148        super().__init__()149        self.attn_dropout = nn.Dropout(attn_dropout)150        self.res_attention = res_attention151        head_dim = d_model // n_heads152        self.scale = nn.Parameter(torch.tensor(head_dim**-0.5), requires_grad=False)153 154    def forward(self, q: Tensor, k: Tensor, v: Tensor, prev: Optional[Tensor] = None):155        attn_scores = torch.matmul(q, k) * self.scale156 157        if prev is not None:158            attn_scores = attn_scores + prev159 160        attn_weights = F.softmax(attn_scores, dim=-1)161        attn_weights = self.attn_dropout(attn_weights)162 163        output = torch.matmul(attn_weights, v)164 165        if self.res_attention:166            return output, attn_weights, attn_scores167        else:168            return output, attn_weights169 170 171class _MultiheadAttention(nn.Module):172    def __init__(173        self,174        d_model,175        n_heads,176        d_k=None,177        d_v=None,178        res_attention=False,179        attn_dropout=0.0,180        proj_dropout=0.0,181        qkv_bias=True,182    ):183        "Multi Head Attention Layer"184 185        super().__init__()186        d_k = d_v = d_model // n_heads187 188        self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v189 190        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)191        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)192        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=qkv_bias)193 194        # Scaled Dot-Product Attention (multiple heads)195        self.res_attention = res_attention196        self.sdp_attn = _ScaledDotProductAttention(197            d_model,198            n_heads,199            attn_dropout=attn_dropout,200            res_attention=self.res_attention,201        )202 203        # Poject output204        self.to_out = nn.Sequential(205            nn.Linear(n_heads * d_v, d_model), nn.Dropout(proj_dropout)206        )207 208    def forward(209        self,210        Q: Tensor,211        K: Optional[Tensor] = None,212        V: Optional[Tensor] = None,213        prev: Optional[Tensor] = None,214    ):215        bs = Q.size(0)216        if K is None:217            K = Q218        if V is None:219            V = Q220 221        # Linear (+ split in multiple heads)222        q_s = (223            self.W_Q(Q).view(bs, -1, self.n_heads, self.d_k).transpose(1, 2)224        )  # q_s: [bs x n_heads x max_q_len x d_k]225        k_s = (226            self.W_K(K).view(bs, -1, self.n_heads, self.d_k).permute(0, 2, 3, 1)227        )  # k_s: [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)228        v_s = (229            self.W_V(V).view(bs, -1, self.n_heads, self.d_v).transpose(1, 2)230        )  # v_s: [bs x n_heads x q_len x d_v]231 232        # Apply Scaled Dot-Product Attention (multiple heads)233        if self.res_attention:234            output, attn_weights, attn_scores = self.sdp_attn(q_s, k_s, v_s, prev=prev)235        else:236            output, attn_weights = self.sdp_attn(q_s, k_s, v_s)237        # output: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len], scores: [bs x n_heads x max_q_len x q_len]238 239        # back to the original inputs dimensions240        output = (241            output.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * self.d_v)242        )  # output: [bs x q_len x n_heads * d_v]243        output = self.to_out(output)244 245        if self.res_attention:246            return output, attn_weights, attn_scores247        else:248            return output, attn_weights249 250 251class Flatten_Head(nn.Module):252    def __init__(self, individual, n_vars, nf, pred_dim):253        super().__init__()254 255        if isinstance(pred_dim, (tuple, list)):256            pred_dim = pred_dim[-1]257        self.individual = individual258        self.n = n_vars if individual else 1259        self.nf, self.pred_dim = nf, pred_dim260 261        if individual:262            self.layers = nn.ModuleList()263            for i in range(self.n):264                self.layers.append(265                    nn.Sequential(nn.Flatten(start_dim=-2), nn.Linear(nf, pred_dim))266                )267        else:268            self.layer = nn.Sequential(269                nn.Flatten(start_dim=-2), nn.Linear(nf, pred_dim)270            )271 272    def forward(self, x: Tensor):273        """274        Args:275            x: [bs x nvars x d_model x n_patch]276            output: [bs x nvars x pred_dim]277        """278        if self.individual:279            x_out = []280            for i, layer in enumerate(self.layers):281                x_out.append(layer(x[:, i]))282            x = torch.stack(x_out, dim=1)283            return x284        else:285            return self.layer(x)286 287 288class _TSTiEncoderLayer(nn.Module):289    def __init__(290        self,291        q_len,292        d_model,293        n_heads,294        d_k=None,295        d_v=None,296        d_ff=256,297        store_attn=False,298        norm="BatchNorm",299        attn_dropout=0,300        dropout=0.0,301        bias=True,302        activation="gelu",303        res_attention=False,304        pre_norm=False,305    ):306        super().__init__()307        assert (308            not d_model % n_heads309        ), f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"310        d_k = d_model // n_heads if d_k is None else d_k311        d_v = d_model // n_heads if d_v is None else d_v312 313        # Multi-Head attention314        self.res_attention = res_attention315        self.self_attn = _MultiheadAttention(316            d_model,317            n_heads,318            d_k,319            d_v,320            attn_dropout=attn_dropout,321            proj_dropout=dropout,322            res_attention=res_attention,323        )324 325        # Add & Norm326        self.dropout_attn = nn.Dropout(dropout)327        if "batch" in norm.lower():328            self.norm_attn = nn.Sequential(329                Transpose(1, 2), nn.BatchNorm1d(d_model), Transpose(1, 2)330            )331        else:332            self.norm_attn = nn.LayerNorm(d_model)333 334        # Position-wise Feed-Forward335        self.ff = nn.Sequential(336            nn.Linear(d_model, d_ff, bias=bias),337            get_act_fn(activation),338            nn.Dropout(dropout),339            nn.Linear(d_ff, d_model, bias=bias),340        )341 342        # Add & Norm343        self.dropout_ffn = nn.Dropout(dropout)344        if "batch" in norm.lower():345            self.norm_ffn = nn.Sequential(346                Transpose(1, 2), nn.BatchNorm1d(d_model), Transpose(1, 2)347            )348        else:349            self.norm_ffn = nn.LayerNorm(d_model)350 351        self.pre_norm = pre_norm352        self.store_attn = store_attn353 354    def forward(self, src: Tensor, prev: Optional[Tensor] = None):355        """356        Args:357            src: [bs x q_len x d_model]358        """359 360        # Multi-Head attention sublayer361        if self.pre_norm:362            src = self.norm_attn(src)363        ## Multi-Head attention364        if self.res_attention:365            src2, attn, scores = self.self_attn(src, src, src, prev)366        else:367            src2, attn = self.self_attn(src, src, src)368        if self.store_attn:369            self.attn = attn370        ## Add & Norm371        src = src + self.dropout_attn(372            src2373        )  # Add: residual connection with residual dropout374        if not self.pre_norm:375            src = self.norm_attn(src)376 377        # Feed-forward sublayer378        if self.pre_norm:379            src = self.norm_ffn(src)380        ## Position-wise Feed-Forward381        src2 = self.ff(src)382        ## Add & Norm383        src = src + self.dropout_ffn(384            src2385        )  # Add: residual connection with residual dropout386        if not self.pre_norm:387            src = self.norm_ffn(src)388 389        if self.res_attention:390            return src, scores391        else:392            return src393 394 395class _TSTiEncoder(nn.Module):  # i means channel-independent396    def __init__(397        self,398        c_in,399        patch_num,400        patch_len,401        n_layers=3,402        d_model=128,403        n_heads=16,404        d_k=None,405        d_v=None,406        d_ff=256,407        norm="BatchNorm",408        attn_dropout=0.0,409        dropout=0.0,410        act="gelu",411        store_attn=False,412        res_attention=True,413        pre_norm=False,414    ):415 416        super().__init__()417 418        self.patch_num = patch_num419        self.patch_len = patch_len420 421        # Input encoding422        q_len = patch_num423        self.W_P = nn.Linear(424            patch_len, d_model425        )  # Eq 1: projection of feature vectors onto a d-dim vector space426        self.seq_len = q_len427 428        # Positional encoding429        W_pos = torch.empty((q_len, d_model))430        nn.init.uniform_(W_pos, -0.02, 0.02)431        self.W_pos = nn.Parameter(W_pos)432 433        # Residual dropout434        self.dropout = nn.Dropout(dropout)435 436        # Encoder437        self.layers = nn.ModuleList(438            [439                _TSTiEncoderLayer(440                    q_len,441                    d_model,442                    n_heads=n_heads,443                    d_k=d_k,444                    d_v=d_v,445                    d_ff=d_ff,446                    norm=norm,447                    attn_dropout=attn_dropout,448                    dropout=dropout,449                    activation=act,450                    res_attention=res_attention,451                    pre_norm=pre_norm,452                    store_attn=store_attn,453                )454                for i in range(n_layers)455            ]456        )457        self.res_attention = res_attention458 459    def forward(self, x: Tensor):460        """461        Args:462            x: [bs x nvars x patch_len x patch_num]463        """464 465        n_vars = x.shape[1]466        # Input encoding467        x = x.permute(0, 1, 3, 2)  # x: [bs x nvars x patch_num x patch_len]468        x = self.W_P(x)  # x: [bs x nvars x patch_num x d_model]469 470        x = torch.reshape(471            x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3])472        )  # x: [bs * nvars x patch_num x d_model]473        x = self.dropout(x + self.W_pos)  # x: [bs * nvars x patch_num x d_model]474 475        # Encoder476        if self.res_attention:477            scores = None478            for mod in self.layers:479                x, scores = mod(x, prev=scores)480        else:481            for mod in self.layers:482                x = mod(x)483        x = torch.reshape(484            x, (-1, n_vars, x.shape[-2], x.shape[-1])485        )  # x: [bs x nvars x patch_num x d_model]486        x = x.permute(0, 1, 3, 2)  # x: [bs x nvars x d_model x patch_num]487 488        return x489 490 491class _PatchTST_backbone(nn.Module):492    def __init__(493        self,494        c_in,495        seq_len,496        pred_dim,497        patch_len,498        stride,499        n_layers=3,500        d_model=128,501        n_heads=16,502        d_k=None,503        d_v=None,504        d_ff=256,505        norm="BatchNorm",506        attn_dropout=0.0,507        dropout=0.0,508        act="gelu",509        res_attention=True,510        pre_norm=False,511        store_attn=False,512        padding_patch=True,513        individual=False,514        revin=True,515        affine=True,516        subtract_last=False,517    ):518 519        super().__init__()520 521        self.revin = revin522        self.revin_layer = RevIN(c_in, affine=affine, subtract_last=subtract_last)523 524        self.patch_len = patch_len525        self.stride = stride526        self.padding_patch = padding_patch527        patch_num = int((seq_len - patch_len) / stride + 1) + 1528        self.patch_num = patch_num529        self.padding_patch_layer = nn.ReplicationPad1d((stride, 0))530 531        self.unfold = nn.Unfold(kernel_size=(1, patch_len), stride=stride)532        self.patch_len = patch_len533 534        self.backbone = _TSTiEncoder(535            c_in,536            patch_num=patch_num,537            patch_len=patch_len,538            n_layers=n_layers,539            d_model=d_model,540            n_heads=n_heads,541            d_k=d_k,542            d_v=d_v,543            d_ff=d_ff,544            attn_dropout=attn_dropout,545            dropout=dropout,546            act=act,547            res_attention=res_attention,548            pre_norm=pre_norm,549            store_attn=store_attn,550        )551 552        # Head553        self.head_nf = d_model * patch_num554        self.n_vars = c_in555        self.individual = individual556        self.head = Flatten_Head(self.individual, self.n_vars, self.head_nf, pred_dim)557 558    def forward(self, z: Tensor):559        """560        Args:561            z: [bs x c_in x seq_len]562        """563 564        if self.revin:565            z = self.revin_layer(z, torch.tensor(True, dtype=torch.bool))566 567        z = self.padding_patch_layer(z)568        b, c, s = z.size()569        z = z.reshape(-1, 1, 1, s)570        z = self.unfold(z)571        z = z.permute(0, 2, 1).reshape(b, c, -1, self.patch_len).permute(0, 1, 3, 2)572 573        z = self.backbone(z)574        z = self.head(z)575 576        if self.revin:577            z = self.revin_layer(z, torch.tensor(False, dtype=torch.bool))578        return z579 580 581class PatchTST(nn.Module, PyTorchModelHubMixin):582    def __init__(583        self,584        c_in,585        c_out,586        seq_len,587        pred_dim=None,588        n_layers=2,589        n_heads=8,590        d_model=512,591        d_ff=2048,592        dropout=0.05,593        attn_dropout=0.0,594        patch_len=16,595        stride=8,596        padding_patch=True,597        revin=True,598        affine=False,599        individual=False,600        subtract_last=False,601        decomposition=False,602        kernel_size=25,603        activation="gelu",604        norm="BatchNorm",605        pre_norm=False,606        res_attention=True,607        store_attn=False,608        classification=False,609    ):610 611        super().__init__()612 613        if pred_dim is None:614            pred_dim = seq_len615 616        self.decomposition = decomposition617        if self.decomposition:618            self.decomp_module = SeriesDecomposition(kernel_size)619            self.model_trend = _PatchTST_backbone(620                c_in=c_in,621                seq_len=seq_len,622                pred_dim=pred_dim,623                patch_len=patch_len,624                stride=stride,625                n_layers=n_layers,626                d_model=d_model,627                n_heads=n_heads,628                d_ff=d_ff,629                norm=norm,630                attn_dropout=attn_dropout,631                dropout=dropout,632                act=activation,633                res_attention=res_attention,634                pre_norm=pre_norm,635                store_attn=store_attn,636                padding_patch=padding_patch,637                individual=individual,638                revin=revin,639                affine=affine,640                subtract_last=subtract_last,641            )642            self.model_res = _PatchTST_backbone(643                c_in=c_in,644                seq_len=seq_len,645                pred_dim=pred_dim,646                patch_len=patch_len,647                stride=stride,648                n_layers=n_layers,649                d_model=d_model,650                n_heads=n_heads,651                d_ff=d_ff,652                norm=norm,653                attn_dropout=attn_dropout,654                dropout=dropout,655                act=activation,656                res_attention=res_attention,657                pre_norm=pre_norm,658                store_attn=store_attn,659                padding_patch=padding_patch,660                individual=individual,661                revin=revin,662                affine=affine,663                subtract_last=subtract_last,664            )665            self.patch_num = self.model_trend.patch_num666        else:667            self.model = _PatchTST_backbone(668                c_in=c_in,669                seq_len=seq_len,670                pred_dim=pred_dim,671                patch_len=patch_len,672                stride=stride,673                n_layers=n_layers,674                d_model=d_model,675                n_heads=n_heads,676                d_ff=d_ff,677                norm=norm,678                attn_dropout=attn_dropout,679                dropout=dropout,680                act=activation,681                res_attention=res_attention,682                pre_norm=pre_norm,683                store_attn=store_attn,684                padding_patch=padding_patch,685                individual=individual,686                revin=revin,687                affine=affine,688                subtract_last=subtract_last,689            )690            self.patch_num = self.model.patch_num691        self.classification = classification692 693    def forward(self, x):694        if self.decomposition:695            res_init, trend_init = self.decomp_module(x)696            res = self.model_res(res_init)697            trend = self.model_trend(trend_init)698            x = res + trend699        else:700            x = self.model(x)701 702        if self.classification:703            x = x.squeeze(-2)704        return x705