File size: 4,557 Bytes
35f60b9 a532d9f 35f60b9 | 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 | import torch
import torch.nn as nn
from functools import partial
from mamba_ssm import Mamba
from models.transformer_modules import Mlp, GatedMLP, Block
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from models.utils import flip_padded_sequences
class MambaBlock(nn.Module):
"""
Wrapper block around the Mamba SSM layer providing the usual
transformer-style residual + MLP structure.
If `mamba-ssm` is not installed, the block falls back to an
identity (no-op) for the SSM portion while still running the
feed-forward MLP path so networks remain usable.
"""
def __init__(self, d_model, d_state=16, d_conv=4, expand=2,
mlp_ratio=4., drop=0.1, drop_path=0.1, act_layer=nn.GELU,
norm_layer=nn.LayerNorm):
super().__init__()
self.norm1 = norm_layer(d_model)
# Mamba SSM layer (state space model)
if Mamba is not None:
self.mamba = Mamba(d_model=d_model, d_state=d_state, d_conv=d_conv, expand=expand)
else:
self.mamba = nn.Identity()
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
# self.norm2 = norm_layer(d_model)
# mlp_hidden_dim = int(d_model * mlp_ratio)
# self.mlp = GatedMLP(in_features=d_model, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
def forward(self, x, mask=None):
if mask is not None:
x = x.masked_fill(~mask.unsqueeze(-1), 0.0)
# x: (B, L, D)
# SSM (Mamba) path with residual
x_in = x
x = self.mamba(self.norm1(x))
if mask is not None:
x = x.masked_fill(~mask.unsqueeze(-1), 0.0)
x = x_in + self.drop_path(x)
return x
class SeqEncoder(nn.Module):
def __init__(self,
seq_dim,
depth=2,
d_state=16,
d_conv=4,
expand=2,
bidi=True,
use_transformer=False,
norm_layer=nn.LayerNorm):
super().__init__()
self.bidi = bidi
self.use_transformer = use_transformer
if use_transformer:
assert not bidi, "Bidirectional mode is not supported with transformer encoder"
# encoder_layer = nn.TransformerEncoderLayer(d_model=seq_dim, nhead=8, dropout=0.1, batch_first=True, norm_first=True)
# self.seq_layers = nn.TransformerEncoder(encoder_layer, num_layers=depth*2)
self.seq_layers = nn.ModuleList([
Block(dim=seq_dim, num_heads=4, drop=0.1, drop_path=0.0, norm_layer=norm_layer)
for i in range(depth*2)
])
else:
self.seq_layers = nn.ModuleList([
MambaBlock(d_model=seq_dim, d_state=d_state, d_conv=d_conv, expand=expand, norm_layer=norm_layer)
for i in range(depth)
])
self.norm_fwd = norm_layer(seq_dim)
# self.post_norm_fwd = norm_layer(seq_dim)
if self.bidi:
self.seq_layers_bwd = nn.ModuleList([
MambaBlock(d_model=seq_dim, d_state=d_state, d_conv=d_conv, expand=expand, norm_layer=norm_layer)
for i in range(depth)
])
self.norm_bwd = norm_layer(seq_dim)
# encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dropout=dropout, batch_first=True, norm_first=True)
# self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
def forward(self, x, mask=None):
if self.bidi:
# Process forward and backward paths independently
xr = flip_padded_sequences(x, mask)
mr = mask
for blk, blk_bwd in zip(self.seq_layers, self.seq_layers_bwd):
xf = blk(x, mask)
xr = blk_bwd(xr, mr)
xf = self.norm_fwd(xf)
xr = self.norm_bwd(xr)
return xf, xr
else:
if self.use_transformer:
mask = ~mask
# # For transformer, we need to provide a key padding mask
# x = self.seq_layers(x, src_key_padding_mask=~mask)
# else:
for blk in self.seq_layers:
x = blk(x, mask)
x = self.norm_fwd(x)
return x |