Update mcdpmamba.py
Browse files- mcdpmamba.py +3 -28
mcdpmamba.py
CHANGED
|
@@ -7,8 +7,6 @@ from einops import rearrange, reduce
|
|
| 7 |
from math import ceil
|
| 8 |
from mamba import Mamba, MambaConfig
|
| 9 |
|
| 10 |
-
|
| 11 |
-
|
| 12 |
class FeedForward(nn.Module):
|
| 13 |
def __init__(self, dim, hidden_dim, dropout):
|
| 14 |
super().__init__()
|
|
@@ -22,11 +20,6 @@ class FeedForward(nn.Module):
|
|
| 22 |
def forward(self, x):
|
| 23 |
return self.net(x)
|
| 24 |
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
class MCGatingUnit(nn.Module):
|
| 31 |
def __init__(self,d_model,d_ffn,dropout):
|
| 32 |
super().__init__()
|
|
@@ -37,8 +30,6 @@ class MCGatingUnit(nn.Module):
|
|
| 37 |
|
| 38 |
self.COB_2 = Mamba(self.config)
|
| 39 |
|
| 40 |
-
|
| 41 |
-
|
| 42 |
def forward(self, x):
|
| 43 |
u, v = x, x
|
| 44 |
u = self.COB_1(u)
|
|
@@ -46,9 +37,8 @@ class MCGatingUnit(nn.Module):
|
|
| 46 |
out = u * v
|
| 47 |
return out
|
| 48 |
|
| 49 |
-
|
| 50 |
class MCDPMAMBABlock(nn.Module):
|
| 51 |
-
def __init__(self, d_model, d_ffn,dropout):
|
| 52 |
super().__init__()
|
| 53 |
|
| 54 |
self.norm = nn.LayerNorm(d_model)
|
|
@@ -65,22 +55,14 @@ class MCDPMAMBABlock(nn.Module):
|
|
| 65 |
out = x + residual
|
| 66 |
return out
|
| 67 |
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
class MCDPMAMBA(nn.Module):
|
| 76 |
def __init__(self, d_model, d_ffn, num_layers, dropout):
|
| 77 |
super().__init__()
|
| 78 |
|
| 79 |
self.model = nn.Sequential(
|
| 80 |
|
| 81 |
-
*[MCDPMAMBABlock(d_model,d_ffn,dropout) for _ in range(num_layers)],
|
| 82 |
-
|
| 83 |
-
|
| 84 |
)
|
| 85 |
|
| 86 |
def forward(self, x):
|
|
@@ -88,10 +70,3 @@ class MCDPMAMBA(nn.Module):
|
|
| 88 |
x = self.model(x)
|
| 89 |
|
| 90 |
return x
|
| 91 |
-
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
|
| 95 |
-
|
| 96 |
-
|
| 97 |
-
|
|
|
|
| 7 |
from math import ceil
|
| 8 |
from mamba import Mamba, MambaConfig
|
| 9 |
|
|
|
|
|
|
|
| 10 |
class FeedForward(nn.Module):
|
| 11 |
def __init__(self, dim, hidden_dim, dropout):
|
| 12 |
super().__init__()
|
|
|
|
| 20 |
def forward(self, x):
|
| 21 |
return self.net(x)
|
| 22 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 23 |
class MCGatingUnit(nn.Module):
|
| 24 |
def __init__(self,d_model,d_ffn,dropout):
|
| 25 |
super().__init__()
|
|
|
|
| 30 |
|
| 31 |
self.COB_2 = Mamba(self.config)
|
| 32 |
|
|
|
|
|
|
|
| 33 |
def forward(self, x):
|
| 34 |
u, v = x, x
|
| 35 |
u = self.COB_1(u)
|
|
|
|
| 37 |
out = u * v
|
| 38 |
return out
|
| 39 |
|
|
|
|
| 40 |
class MCDPMAMBABlock(nn.Module):
|
| 41 |
+
def __init__(self, d_model, d_ffn, dropout):
|
| 42 |
super().__init__()
|
| 43 |
|
| 44 |
self.norm = nn.LayerNorm(d_model)
|
|
|
|
| 55 |
out = x + residual
|
| 56 |
return out
|
| 57 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 58 |
class MCDPMAMBA(nn.Module):
|
| 59 |
def __init__(self, d_model, d_ffn, num_layers, dropout):
|
| 60 |
super().__init__()
|
| 61 |
|
| 62 |
self.model = nn.Sequential(
|
| 63 |
|
| 64 |
+
*[MCDPMAMBABlock(d_model, d_ffn, dropout) for _ in range(num_layers)],
|
| 65 |
+
|
|
|
|
| 66 |
)
|
| 67 |
|
| 68 |
def forward(self, x):
|
|
|
|
| 70 |
x = self.model(x)
|
| 71 |
|
| 72 |
return x
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|