| 16 | |
| 17 | # Cell |
| 18 | class PatchTST_backbone(nn.Module): |
| 19 | def __init__(self, c_in:int, context_window:int, target_window:int, patch_len:int, stride:int, max_seq_len:Optional[int]=1024, |
| 20 | n_layers:int=3, d_model=128, n_heads=16, d_k:Optional[int]=None, d_v:Optional[int]=None, |
| 21 | d_ff:int=256, norm:str='BatchNorm', attn_dropout:float=0., dropout:float=0., act:str="gelu", key_padding_mask:bool='auto', |
| 22 | padding_var:Optional[int]=None, attn_mask:Optional[Tensor]=None, res_attention:bool=True, pre_norm:bool=False, store_attn:bool=False, |
| 23 | pe:str='zeros', learn_pe:bool=True, fc_dropout:float=0., head_dropout = 0, padding_patch = None, |
| 24 | pretrain_head:bool=False, head_type = 'flatten', individual = False, revin = True, affine = True, subtract_last = False, |
| 25 | verbose:bool=False, **kwargs): |
| 26 | |
| 27 | super().__init__() |
| 28 | |
| 29 | # RevIn |
| 30 | self.revin = revin |
| 31 | if self.revin: self.revin_layer = RevIN(c_in, affine=affine, subtract_last=subtract_last) |
| 32 | |
| 33 | # Patching |
| 34 | self.patch_len = patch_len |
| 35 | self.stride = stride |
| 36 | self.padding_patch = padding_patch |
| 37 | patch_num = int((context_window - patch_len)/stride + 1) |
| 38 | if padding_patch == 'end': # can be modified to general case |
| 39 | self.padding_patch_layer = nn.ReplicationPad1d((0, stride)) |
| 40 | patch_num += 1 |
| 41 | |
| 42 | # Backbone |
| 43 | self.backbone = TSTiEncoder(c_in, patch_num=patch_num, patch_len=patch_len, max_seq_len=max_seq_len, |
| 44 | n_layers=n_layers, d_model=d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, |
| 45 | attn_dropout=attn_dropout, dropout=dropout, act=act, key_padding_mask=key_padding_mask, padding_var=padding_var, |
| 46 | attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm, store_attn=store_attn, |
| 47 | pe=pe, learn_pe=learn_pe, verbose=verbose, **kwargs) |
| 48 | |
| 49 | # Head |
| 50 | self.head_nf = d_model * patch_num |
| 51 | self.n_vars = c_in |
| 52 | self.pretrain_head = pretrain_head |
| 53 | self.head_type = head_type |
| 54 | self.individual = individual |
| 55 | |
| 56 | if self.pretrain_head: |
| 57 | self.head = self.create_pretrain_head(self.head_nf, c_in, fc_dropout) # custom head passed as a partial func with all its kwargs |
| 58 | elif head_type == 'flatten': |
| 59 | self.head = Flatten_Head(self.individual, self.n_vars, self.head_nf, target_window, head_dropout=head_dropout) |
| 60 | |
| 61 | |
| 62 | def forward(self, z, edge_index, edge_attr): # z: [bs x nvars x seq_len] |
| 63 | # norm |
| 64 | if self.revin: |
| 65 | z = z.permute(0,2,1) |
| 66 | z = self.revin_layer(z, 'norm') |
| 67 | z = z.permute(0,2,1) |
| 68 | |
| 69 | # do patching |
| 70 | if self.padding_patch == 'end': |
| 71 | z = self.padding_patch_layer(z) |
| 72 | z = z.unfold(dimension=-1, size=self.patch_len, step=self.stride) # z: [bs x nvars x patch_num x patch_len] |
| 73 | z = z.permute(0,1,3,2) # z: [bs x nvars x patch_len x patch_num] |
| 74 | # print(z.shape,'z shape ...') |
| 75 | # model |