(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)
| 208 | ### Add DG-Conv and SG-Conv |
| 209 | class TSTEncoderLayer(nn.Module): |
| 210 | def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=256, store_attn=False, |
| 211 | norm='BatchNorm', attn_dropout=0, dropout=0., bias=True, activation="gelu", res_attention=False, pre_norm=False): |
| 212 | super().__init__() |
| 213 | assert not d_model%n_heads, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})" |
| 214 | d_k = d_model // n_heads if d_k is None else d_k |
| 215 | d_v = d_model // n_heads if d_v is None else d_v |
| 216 | self.d_model = d_model |
| 217 | # Graph-Conv |
| 218 | self.g_conv = G_Conv(q_len*d_model, q_len*d_model,K=3) |
| 219 | self.g_dy_conv = G_Dy_Conv(q_len*d_model, q_len*d_model,K=3) |
| 220 | |
| 221 | # Multi-Head attention |
| 222 | self.res_attention = res_attention |
| 223 | self.self_attn = _MultiheadAttention(d_model, n_heads, d_k, d_v, attn_dropout=attn_dropout, proj_dropout=dropout, res_attention=res_attention) |
| 224 | |
| 225 | # Add & Norm |
| 226 | self.dropout_attn = nn.Dropout(dropout) |
| 227 | if "batch" in norm.lower(): |
| 228 | self.norm_attn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2)) |
| 229 | else: |
| 230 | self.norm_attn = nn.LayerNorm(d_model) |
| 231 | |
| 232 | # Position-wise Feed-Forward |
| 233 | self.ff = nn.Sequential(nn.Linear(d_model, d_ff, bias=bias), |
| 234 | get_activation_fn(activation), |
| 235 | nn.Dropout(dropout), |
| 236 | nn.Linear(d_ff, d_model, bias=bias)) |
| 237 | |
| 238 | # Add & Norm |
| 239 | self.dropout_ffn = nn.Dropout(dropout) |
| 240 | if "batch" in norm.lower(): |
| 241 | self.norm_ffn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2)) |
| 242 | else: |
| 243 | self.norm_ffn = nn.LayerNorm(d_model) |
| 244 | |
| 245 | self.pre_norm = pre_norm |
| 246 | self.store_attn = store_attn |
| 247 | |
| 248 | |
| 249 | def forward(self, src:Tensor, |
nothing calls this directly
no test coverage detected