KeondoPark

Code refactoring WIP

a532d9f 5 months ago

4.56 kB

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

	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