(self, z, edge_index, edge_attr)
| 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 |
| 76 | z = self.backbone(z, edge_index, edge_attr) # z: [bs x nvars x d_model x patch_num] |
| 77 | z = self.head(z) # z: [bs x nvars x target_window] |
| 78 | |
| 79 | # denorm |
| 80 | if self.revin: |
| 81 | z = z.permute(0,2,1) |
| 82 | z = self.revin_layer(z, 'denorm') |
| 83 | z = z.permute(0,2,1) |
| 84 | return z |
| 85 | |
| 86 | def create_pretrain_head(self, head_nf, vars, dropout): |
| 87 | return nn.Sequential(nn.Dropout(dropout), |
nothing calls this directly
no outgoing calls
no test coverage detected