| 179 | |
| 180 | # Cell |
| 181 | class TSTEncoder(nn.Module): |
| 182 | def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=None, |
| 183 | norm='BatchNorm', attn_dropout=0., dropout=0., activation='gelu', |
| 184 | res_attention=False, n_layers=1, pre_norm=False, store_attn=False): |
| 185 | super().__init__() |
| 186 | |
| 187 | self.layers = nn.ModuleList([TSTEncoderLayer(q_len, d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm, |
| 188 | attn_dropout=attn_dropout, dropout=dropout, |
| 189 | activation=activation, res_attention=res_attention, |
| 190 | pre_norm=pre_norm, store_attn=store_attn) for i in range(n_layers)]) |
| 191 | self.res_attention = res_attention |
| 192 | |
| 193 | def forward(self, src:Tensor, |
| 194 | edge_index: torch.LongTensor, |
| 195 | edge_weight: torch.FloatTensor, |
| 196 | key_padding_mask:Optional[Tensor]=None, |
| 197 | attn_mask:Optional[Tensor]=None): |
| 198 | output = src |
| 199 | scores = None |
| 200 | if self.res_attention: |
| 201 | for mod in self.layers: output, scores = mod(output,edge_index, edge_weight, prev=scores, key_padding_mask=key_padding_mask, attn_mask=attn_mask) |
| 202 | return output |
| 203 | else: |
| 204 | for mod in self.layers: output = mod(output, key_padding_mask=key_padding_mask, attn_mask=attn_mask) |
| 205 | return output |
| 206 | |
| 207 | |
| 208 | ### Add DG-Conv and SG-Conv |