File size: 4,812 Bytes
35f60b9 a532d9f 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 | import torch
import torch.nn as nn
from models.transformer_modules import BlockFlashAttn, Mlp, freq_pos_enc, PatchUpsample
class FusionModule(nn.Module):
def __init__(self, encoder_embed_dim=128,
num_heads=4, mlp_ratio=4, norm_layer=nn.LayerNorm, fusion_depth=2,
fusion_method='concat', pos_enc='sinusoidal'):
super().__init__()
self.fusion_method = fusion_method
self.num_heads = num_heads
if self.fusion_method == 'concat':
self.prenorm = norm_layer(encoder_embed_dim * 2) # prenorm for the weighted sum features
# self.merge = nn.Linear(encoder_embed_dim * 2, encoder_embed_dim)
self.merge = Mlp(in_features=encoder_embed_dim * 2, out_features=encoder_embed_dim)
elif self.fusion_method == 'cross_attn':
self.cross_attention_blocks = nn.ModuleList([
BlockFlashAttn(encoder_embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
# BlockFlashAttn(encoder_embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
for i in range(fusion_depth)])
self.post_norm = norm_layer(encoder_embed_dim)
self.pos_enc = pos_enc
if pos_enc == 'rope':
self.rope = freq_pos_enc(dim=encoder_embed_dim // num_heads, max_seq_len=128)
def forward(self, x, y):
if self.fusion_method == 'concat':
x = torch.cat([x, y], dim=-1)
x = self.prenorm(x) # apply prenorm before the forward_features
x = self.merge(x)
elif self.fusion_method == 'cross_attn':
if self.pos_enc == 'rope':
pe = self.rope.unsqueeze(0).to(dtype=x.dtype, device=x.device) # (1, max_seq_len, dim)
pe = pe[:, :x.shape[1], :] # (1, seq_len, dim)
pe = pe.repeat(x.shape[0], 1, 1) # (N, max_seq_len, dim)
pe2 = None
if not x.shape[1] == y.shape[1]:
pe2 = freq_pos_enc(dim=y.shape[2] // self.num_heads, max_seq_len=y.shape[1])
pe2 = pe2.unsqueeze(0).to(dtype=x.dtype, device=x.device) # (1, max_seq_len, dim)
pe2 = pe2.repeat(x.shape[0], 1, 1) # (N, max_seq_len, dim)
else:
pe = None
for blk in self.cross_attention_blocks:
x, _ = blk(x, y, rope=pe, rope2=pe2)
x = self.post_norm(x)
return x
class FusionUpsampleBlock(nn.Module):
def __init__(self, dim_list, fusion_method='cross_attn', fusion_depth=2, fusion_heads=4,
mlp_ratio=4., norm_layer=nn.LayerNorm, pos_enc='sinusoidal',
use_shared=True, use_upsample=True, modalities=["EEG","EOG"]):
super().__init__()
self.fusion_method = fusion_method
self.use_shared = use_shared
self.use_upsample = use_upsample
self.fusion = nn.ModuleDict()
self.upsampler = nn.ModuleDict()
self.dim_list = dim_list
self.modalities = modalities
for m in self.modalities:
if use_shared:
self.fusion[m] = FusionModule(encoder_embed_dim=dim_list[-1],
mlp_ratio=mlp_ratio, norm_layer=norm_layer,
fusion_method=fusion_method, fusion_depth=fusion_depth, num_heads=fusion_heads, pos_enc=pos_enc)
if use_upsample:
upsampler_list = []
for i in range(len(dim_list) - 1):
up = 2 if i > 0 else 5
upsampler_list.append(PatchUpsample(embed_dim=dim_list[1+i], up_dim=dim_list[i], upsample=up, norm_layer=norm_layer))
self.upsampler[m] = nn.ModuleList(upsampler_list)
def forward(self, x_shared, x_specific):
latents_fused = []
upsampled = []
for i, m in enumerate(self.modalities):
if self.use_shared:
x_fused = self.fusion[m](x_specific[-1][i], x_shared[-1][i])
else:
x_fused = x_specific[-1][i]
latents_fused.append(x_fused)
if self.use_upsample:
for j in range(len(self.dim_list)-2, -1, -1):
x_fused = self.upsampler[m][j](x_fused, x_specific[j][i])
upsampled.append(x_fused)
return latents_fused, upsampled
|