File size: 4,346 Bytes
493bf96 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import concurrent.futures
import threading
import torch.nn as nn
import torch.nn.functional as F
import torch
from mamba_ssm import Mamba
from layers.SelfAttention_Family import FullAttention, AttentionLayer
import torch
import torch.nn as nn
class FeedForwardNetwork(nn.Module):
def __init__(self, nvars, dmodel, dff, drop=0.1):
super(FeedForwardNetwork, self).__init__()
self.ffn1pw1 = nn.Conv1d(in_channels=nvars * dmodel, out_channels=nvars * dff,
kernel_size=1, stride=1, padding=0, dilation=1, groups=nvars)
self.ffn1act = nn.GELU()
self.ffn1pw2 = nn.Conv1d(in_channels=nvars * dff, out_channels=nvars * dmodel,
kernel_size=1, stride=1, padding=0, dilation=1, groups=nvars)
self.ffn1drop1 = nn.Dropout(drop)
self.ffn1drop2 = nn.Dropout(drop)
def forward(self, x):
x = self.ffn1drop1(self.ffn1pw1(x))
x = self.ffn1act(x)
x = self.ffn1drop2(self.ffn1pw2(x))
return x
class EncoderLayer(nn.Module):
def __init__(self, attention, attention_r, d_model, d_ff=None, dropout=0.1, activation="relu"):
super(EncoderLayer, self).__init__()
d_ff = d_ff or 4 * d_model
self.attention = attention
self.attention_r = attention_r
self.ffn1 = FeedForwardNetwork(nvars=1, dmodel=d_model, dff=d_ff)
self.ffn2 = FeedForwardNetwork(nvars=1, dmodel=d_ff, dff=d_model)
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
self.activation = F.relu if activation == "relu" else F.gelu
self.man = Mamba(
d_model=11, # Model dimension d_model
d_state=16, # SSM state expansion factor
d_conv=2, # Local convolution width
expand=1, # Block expansion factor)
)
self.man2 = Mamba(
d_model=11, # Model dimension d_model
d_state=16, # SSM state expansion factor
d_conv=2, # Local convolution width
expand=1, # Block expansion factor)
)
self.a = AttentionLayer(
FullAttention(False, 2, attention_dropout=0.1,
output_attention=True), 11,1)
def forward(self, x, attn_mask=None, tau=None, delta=None):
x = x.permute(0, 2, 1)
x = self.ffn1(x)
x = self.ffn2(x)
x = x.permute(0, 2, 1)
new_x = self.attention(x) + self.attention_r(x.flip(dims=[1])).flip(dims=[1])
attn = 1
x = x + new_x #torch.Size([32, 11, 256])
# print(x.shape)
y = x = self.norm1(x)
# print(y.shape)
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
y = self.dropout(self.conv2(y).transpose(-1, 1))
return self.norm2(x + y), attn
class Encoder(nn.Module):
def __init__(self, attn_layers, conv_layers=None, norm_layer=None):
super(Encoder, self).__init__()
self.attn_layers = nn.ModuleList(attn_layers)
self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
self.norm = norm_layer
def forward(self, x, attn_mask=None, tau=None, delta=None):
# x [B, L, D]
# print(x.shape)
attns = []
if self.conv_layers is not None:
for i, (attn_layer, conv_layer) in enumerate(zip(self.attn_layers, self.conv_layers)):
delta = delta if i == 0 else None
x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta)
x = conv_layer(x)
attns.append(attn)
x, attn = self.attn_layers[-1](x, tau=tau, delta=None)
attns.append(attn)
else:
for attn_layer in self.attn_layers:
x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta)
attns.append(attn)
if self.norm is not None:
x = self.norm(x)
return x, attns
|