mirror/TSLANet
2
0
Fork 0
mirror of https://github.com/emadeldeen24/TSLANet.git synced 2024-11-21 00:40:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update TSLANet_Forecasting.py

Browse Source
This commit is contained in:
Emadeldeen Eldele 2024-04-19 09:28:47 +08:00 committed by GitHub
parent cc528c563d
commit f00454cc16
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
Forecasting/TSLANet_Forecasting.py
View File

@ -108,25 +108,26 @@ class Adaptive_Spectral_Block(nn.Module):
class TSLANet_layer(L.LightningModule):
def __init__(self, dim, mlp_ratio=3., drop=0., drop_path=0., norm_layer=nn.LayerNorm):
super().__init__()
self.norm1 = norm_layer(dim)
self.filter = Adaptive_Spectral_Block(dim)
self.asb = Adaptive_Spectral_Block(dim)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.icb = ICB(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop)
def forward(self, x):
if not args.ICB:
# Check if both ASB and ICB are true
if args.ICB and args.ASB:
x = x + self.drop_path(self.icb(self.norm2(self.asb(self.norm1(x)))))
# If only ICB is true
elif args.ICB:
x = x + self.drop_path(self.icb(self.norm2(x)))
return x
if not args.ASB:
x = x + self.drop_path(self.filter(self.norm1(x)))
return x
x = x + self.drop_path(self.icb(self.norm2(self.filter(self.norm1(x)))))
# If only ASB is true
elif args.ASB:
x = x + self.drop_path(self.asb(self.norm1(x)))
# If neither is true, just pass x through
return x