| 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) |
| |
| 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) |
| |
| 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) |
| 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) |
| pe = pe[:, :x.shape[1], :] |
| pe = pe.repeat(x.shape[0], 1, 1) |
| 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) |
| pe2 = pe2.repeat(x.shape[0], 1, 1) |
| |
| 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 |
| |
