Class TSTEncoderLayer

LDPS_Graph/layers/PatchTST_backbone.py:209–298

### Add DG-Conv and SG-Conv
class TSTEncoderLayer(nn.Module):
    def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=256, store_attn=False,
                 norm='BatchNorm', attn_dropout=0, dropout=0., bias=True, activation="gelu", res_attention=False, pre_norm=False):
        super().__init__()
        assert not d_model%n_heads, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
        d_k = d_model // n_heads if d_k is None else d_k
        d_v = d_model // n_heads if d_v is None else d_v
        self.d_model = d_model
        # Graph-Conv
        self.g_conv = G_Conv(q_len*d_model, q_len*d_model,K=3)
        self.g_dy_conv = G_Dy_Conv(q_len*d_model, q_len*d_model,K=3)

        # Multi-Head attention
        self.res_attention = res_attention
        self.self_attn = _MultiheadAttention(d_model, n_heads, d_k, d_v, attn_dropout=attn_dropout, proj_dropout=dropout, res_attention=res_attention)

        # Add & Norm
        self.dropout_attn = nn.Dropout(dropout)
        if "batch" in norm.lower():
            self.norm_attn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))
        else:
            self.norm_attn = nn.LayerNorm(d_model)

        # Position-wise Feed-Forward
        self.ff = nn.Sequential(nn.Linear(d_model, d_ff, bias=bias),
                                get_activation_fn(activation),
                                nn.Dropout(dropout),
                                nn.Linear(d_ff, d_model, bias=bias))

        # Add & Norm
        self.dropout_ffn = nn.Dropout(dropout)
        if "batch" in norm.lower():
            self.norm_ffn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))
        else:
            self.norm_ffn = nn.LayerNorm(d_model)

        self.pre_norm = pre_norm
        self.store_attn = store_attn


    def forward(self, src:Tensor, 
                edge_index: torch.LongTensor,
                edge_weight: torch.FloatTensor,
                prev:Optional[Tensor]=None,
                 key_padding_mask:Optional[Tensor]=None, 
                 attn_mask:Optional[Tensor]=None) -> Tensor:

        # Multi-Head attention sublayer
        if self.pre_norm:
            src = self.norm_attn(src)

        # print(src.shape,'src shape 111')
        # G-input
        # g_in_src = src.reshape(src.size(0),-1)
        
        # # # G-Conv
        # g_src = self.g_conv(g_in_src, edge_index, edge_weight)
        # g_src = g_src.view(g_in_src.size(0), -1, self.d_model)