from functools import partial import copy import torch import torch.nn as nn import torch.nn.functional as F import numpy as np from models.encoders import MultiEncoderClassification from models.transformer_modules import RevIN from models.mamba_modules import SeqEncoder from models.fusion import FusionUpsampleBlock from models.utils import flip_padded_sequences from einops import rearrange from configs.global_cfg import TARGET_RESAMPLE class ClassificationMultiConv1d(nn.Module): """ Masked Autoencoder with VisionTransformer backbone """ def __init__(self, sig_len=6000, patch_size=10, embed_dim=[32,64,128], depth=[2,2,4], specific_depth=[2,2], num_heads=8, fusion_heads=8, fusion_depth=2, mlp_ratio=4, norm_layer=nn.LayerNorm, num_classes=5, modalities=["EEG","EOG"], num_experts=2, num_experts_per=1, use_shared=True, use_moe=True, fusion_method='concat', use_revin=False, pos_enc='sinusoidal', ft_task='sleep_staging'): super().__init__() self.sig_len = sig_len self.patch_size = patch_size self.modalities = modalities self.num_experts = num_experts self.num_experts_per = num_experts_per self.use_shared = use_shared self.use_moe = use_moe self.fusion_method = fusion_method self.embed_dim = embed_dim self.depth = depth self.pos_enc = pos_enc self.encoder = MultiEncoderClassification(sig_len=sig_len, patch_size=patch_size, embed_dim=embed_dim, depth=depth, num_heads=num_heads, specific_depth=specific_depth, mlp_ratio=mlp_ratio, norm_layer=norm_layer, modalities=self.modalities, num_experts=num_experts, num_experts_per=num_experts_per, use_shared=use_shared, use_moe=use_moe, pos_enc=pos_enc) # self.head = nn.Linear(embed_dim[-1], num_classes) if num_classes > 0 else nn.Identity() self.num_patches = self.encoder.num_patches self.fuser = FusionUpsampleBlock(dim_list=embed_dim[2:], fusion_method=fusion_method, fusion_depth=fusion_depth, fusion_heads=fusion_heads, mlp_ratio=mlp_ratio, norm_layer=norm_layer, pos_enc=pos_enc, use_shared=use_shared, use_upsample=True, # if ft_task=='apnea_seg' else False, modalities=self.modalities) self.ft_task = ft_task if self.ft_task == 'sleep_staging': post_fusion_dim = embed_dim[3] * len(self.modalities) self.head = nn.Linear(post_fusion_dim, num_classes) if num_classes > 0 else nn.Identity() elif self.ft_task == 'apnea_seg': post_fusion_dim = embed_dim[2] * len(self.modalities) self.patch_per_sec = TARGET_RESAMPLE // (self.patch_size * 2) # Encoder downsamples by 2 for consistency in pretrainng and finetuning self.head = nn.Linear(post_fusion_dim, num_classes) # 1D conv for sequence labeling self.use_revin = use_revin if self.use_revin: self.revin = nn.ModuleDict() for m in self.modalities: if m != 'SaO2': self.revin[m] = RevIN(num_features=1, affine=False) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): # we use xavier_uniform following official JAX ViT: torch.nn.init.xavier_uniform_(m.weight) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def forward_features(self, inputs): modalities_to_use = sorted(list(inputs.keys())) if self.use_revin: for i, m in enumerate(modalities_to_use): if m == 'SaO2': inputs[m] = inputs[m][:, ::TARGET_RESAMPLE, :] # N,L* TARGET_RESAMPLE,1 -> N,L,1 else: inputs[m] = self.revin[m](inputs[m], mode='norm') latents_shared, latents_specific, gating_scores = self.encoder(inputs, modalities_to_use) latents_fused, upsampled = self.fuser(latents_shared, latents_specific) if self.ft_task == 'apnea_seg': latents = torch.stack(upsampled, dim=0) # (M, N, L, D * 2 or 4) latents = rearrange(latents, 'm n (l s) d -> n l s (m d)', s=self.patch_per_sec) latents = latents.mean(2) lengths = latents.shape[1] if lengths > 30: shift = (lengths - 30) // 2 latents = latents[:, shift:shift+30, :] else: # latents = torch.concat(latents_cat, dim=-1) # (M, N, L, D * 2 or 4) latents = torch.stack(latents_fused, dim=0) latents = rearrange(latents, 'm n l d -> n m l d') latents = latents.mean(2) latents = rearrange(latents, 'n m d -> n (m d)') return latents def forward(self, sig): x = self.forward_features(sig) x = self.head(x) return x class ClassificationMultiConv1dMamba(nn.Module): """ Masked Autoencoder with VisionTransformer backbone """ def __init__(self, sig_len=6000, patch_size=10, embed_dim=[32,64,128], depth=[2,2,4], specific_depth=[2,2], num_heads=8, seq_dim=512, seq_depth=2, mlp_ratio=4, norm_layer=nn.LayerNorm, num_classes=5, modalities=["EEG","EOG"], num_experts=2, num_experts_per=1, use_shared=True, use_moe=True, fusion_method='concat', fusion_depth=2, fusion_heads=8, use_revin=False, pos_enc='sinusoidal', ft_task='sleep_staging', d_state=16, d_conv=4, expand=2, bidi=True, pretrain=False, use_micro=False, demo_keys=['age','bmi','sex'], cycle_len=180, use_lstm=False, use_transformer=False, simclr=False, simclr_temp=0.1, dino=False, dino_student_temp=0.1, dino_teacher_temp=0.04, dino_center_momentum=0.9, dino_momentum=0.996, ): super().__init__() # Consolidate simple config assignments to reduce repetition _cfg = dict( sig_len=sig_len, patch_size=patch_size, modalities=list(modalities), num_experts=num_experts, num_experts_per=num_experts_per, use_shared=use_shared, use_moe=use_moe, fusion_method=fusion_method, embed_dim=embed_dim, pos_enc=pos_enc, pretrain=pretrain, ft_task=ft_task, use_micro=use_micro, demo_keys=list(demo_keys), use_lstm=use_lstm, cycle_len=cycle_len, use_transformer=use_transformer, simclr=simclr, simclr_temp=simclr_temp, dino=dino, dino_student_temp=dino_student_temp, dino_teacher_temp=dino_teacher_temp, dino_center_momentum=dino_center_momentum, dino_momentum=dino_momentum, ) for _k, _v in _cfg.items(): setattr(self, _k, _v) # Build core modules self.encoder = MultiEncoderClassification(sig_len=sig_len, patch_size=patch_size, embed_dim=embed_dim, depth=depth, num_heads=num_heads, specific_depth=specific_depth, mlp_ratio=mlp_ratio, norm_layer=norm_layer, modalities=self.modalities, num_experts=num_experts, num_experts_per=num_experts_per, use_shared=use_shared, use_moe=use_moe, pos_enc=pos_enc) self.num_patches = self.encoder.num_patches self.fuser = FusionUpsampleBlock(dim_list=embed_dim[2:], fusion_method=fusion_method, fusion_depth=fusion_depth, fusion_heads=fusion_heads, mlp_ratio=mlp_ratio, norm_layer=norm_layer, pos_enc=pos_enc, use_shared=use_shared, use_upsample=True, modalities=self.modalities) # ReVin (optional) self.use_revin = use_revin if self.use_revin: self.revin = nn.ModuleDict() for m in self.modalities: if m != 'SaO2': self.revin[m] = RevIN(num_features=1, affine=False) # Heads and sequence modules self._build_heads(embed_dim, num_classes, bidi, seq_dim, seq_depth, d_state, d_conv, expand, use_micro, use_lstm, use_transformer, norm_layer) self.apply(self._init_weights) if self.pretrain and self.dino: self.teacher_encoder = None self.teacher_fuser = None self.teacher_proj = None self.teacher_sequence_encoder = None def _init_weights(self, m): if isinstance(m, nn.Linear): torch.nn.init.xavier_uniform_(m.weight) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) # --- Helper builders ------------------------------------------------- def _build_heads(self, embed_dim, num_classes, bidi, seq_dim, seq_depth, d_state, d_conv, expand, use_micro, use_lstm, use_transformer, norm_layer=nn.LayerNorm): # micro head if use_micro: self.micro_head = nn.Linear(embed_dim[-1] * len(self.modalities), num_classes) return # sequence / macro heads self.bidi = bidi self.proj = nn.Linear(embed_dim[-1] * len(self.modalities), seq_dim) self.sequence_encoder = SeqEncoder(seq_dim=seq_dim, depth=seq_depth, d_state=d_state, d_conv=d_conv, expand=expand, bidi=bidi, use_transformer=use_transformer, norm_layer=norm_layer) out_dim = seq_dim * 2 if self.bidi else seq_dim out_dim += embed_dim[-1] * len(self.modalities) if self.pretrain: # Projection head for contrastive learning self.contr_head = nn.Sequential( nn.Linear(seq_dim, seq_dim), nn.ReLU(), nn.Linear(seq_dim, seq_dim) ) self.koleo_weight = 0.01 self.tau = 0.1 self.upsilon = 0.5 else: if use_lstm: if self.ft_task == 'diagnosis': self.head = DiagnosisFinetuneFullLSTMCOXPHWithDemo(input_dim=out_dim, embed_dim=128, num_classes=num_classes) else: self.head = SleepEventLSTMClassifier(input_dim=out_dim, embed_dim=128, num_classes=num_classes, num_heads=4) else: self.head = nn.Linear(out_dim, num_classes) def forward_features(self, inputs, mask=None): modalities_to_use = sorted(list(inputs.keys())) m = modalities_to_use[0] N, L, _ = inputs[m].shape if mask is not None: nonmask_idx = torch.nonzero(mask.view(-1,)).squeeze(-1) # N*L_masked for i, m in enumerate(modalities_to_use): inputs[m] = inputs[m].view(-1, self.sig_len, inputs[m].shape[-1]) # add channel dim if self.use_revin: if m == 'SaO2': inputs[m] = inputs[m][:, ::TARGET_RESAMPLE, :] # N,L* TARGET_RESAMPLE,1 -> N,L,1 else: inputs[m] = self.revin[m](inputs[m], mode='norm') if mask is not None: inputs[m] = torch.gather(inputs[m], dim=0, index=nonmask_idx.unsqueeze(-1).unsqueeze(-1).repeat(1, inputs[m].shape[1], inputs[m].shape[2])) latents_shared, latents_specific, _ = self.encoder(inputs, modalities_to_use) latents_fused, upsampled = self.fuser(latents_shared, latents_specific) latents = torch.stack(latents_fused, dim=0) # (M, N, L, D) latents = rearrange(latents, 'm n l d -> n l (m d)') latents_micro = latents # latents_micro = torch.stack(latents_fused, dim=0) # (M, N, L, D) # latents_micro = rearrange(latents_micro, 'm n l d -> n l (m d)') if not self.pretrain: latents_micro = latents_micro.mean(1) dummy2 = torch.zeros((N*(L // self.sig_len), latents_micro.shape[-1]), device=latents_micro.device, dtype=latents_micro.dtype) dummy2.scatter_(dim=0, index=nonmask_idx.unsqueeze(-1).repeat(1, latents_micro.shape[-1]), src=latents_micro) latents_micro = dummy2 latents_micro = rearrange(latents_micro, '(n1 n2) d -> n1 n2 d', n1=N) # Finetuning with micro head only (ablation) if self.use_micro: return None, latents_micro else: latents = self.proj(latents) latents = latents.mean(1) dummy = torch.zeros((N*(L // self.sig_len), latents.shape[-1]), device=latents.device, dtype=latents.dtype) dummy.scatter_(dim=0, index=nonmask_idx.unsqueeze(-1).repeat(1, latents.shape[-1]), src=latents) latents = dummy latents = rearrange(latents, '(n1 n2) d -> n1 n2 d', n1=N) latents = self.sequence_encoder(latents, mask=mask) return latents, latents_micro def forward(self, sig, mask=None, demo=None): if self.pretrain and self.dino: return self._forward_dino(sig, mask) if self.pretrain and self.simclr: return self._forward_simclr(sig, mask) feats = self.forward_features(sig, mask=mask) x_macro, x_micro = feats # Pretraining if self.pretrain: if self.bidi: xf, xr = x_macro loss_fwd = self.dgcl_loss(xf, demo, mask) loss_bwd = self.dgcl_loss(xr, demo, mask) loss = (loss_fwd + loss_bwd) / 2 else: loss = self.dgcl_loss(x_macro, demo, mask) return loss # Inference / Finetuning if not self.use_micro and self.bidi: xf, xr = x_macro xr = flip_padded_sequences(xr, mask) x_macro = torch.cat([xf, xr], dim=-1) B = x_micro.shape[0] lengths = mask.sum(dim=1).long() # (B,) if self.use_micro: if self.ft_task in ['diagnosis', 'age', 'sex', 'ahi']: x_micro = torch.stack([x_micro[i, :lengths[i]].mean(dim=0) for i in range(B)]) x = self.micro_head(x_micro) else: x_mami = torch.cat([x_macro, x_micro], dim=-1) if self.use_lstm: x = self.head(x_mami, mask, demo) else: if self.ft_task in ['diagnosis', 'age', 'sex', 'ahi']: x_mami = torch.stack([x_mami[i, :lengths[i]].mean(dim=0) for i in range(B)]) x = self.head(x_mami) return x # if self.ft_task in ['diagnosis', 'age' ,'sex', 'ahi']: # B = x_micro.shape[0] # lengths = mask.sum(dim=1).long() # (B,) # if self.use_micro: # x_micro = torch.stack([x_micro[i, :lengths[i]].mean(dim=0) for i in range(B)]) # x = self.micro_head(x_micro) # else: # x_mami = torch.cat([x_macro, x_micro], dim=-1) # if self.use_lstm: # x = self.head(x_mami, mask, demo) # else: # x_mami = torch.stack([x_mami[i, :lengths[i]].mean(dim=0) for i in range(B)]) # x = self.head(x_mami) # return x # else: # if self.use_micro: # x = self.micro_head(x_micro) # else: # x_mami = torch.cat([x_macro, x_micro], dim=-1) # # x = self.head(x_mami) # if self.use_lstm: # x = self.head(x_mami, mask) # else: # x = self.head(x_mami) # x = rearrange(x, 'n1 n2 d -> n1 d n2') # return x def dgcl_loss(self, xf, demo, mask): """ Calculate DGCL loss based on the difference of age. When the corresponding ages between two sequences are distant, the representation should be distant In constrast, if ages are close, the representation should be similar. """ N, L, D = xf.shape num_cycle = L // self.cycle_len age = demo['age'] sex = demo['sex'] bmi = demo['bmi'] mu_age, std_age = 56.071731, 16.131597 mu_bmi, std_bmi = 27.095034, 4.943150 age = (age - mu_age) / std_age # Normalize age_t = age.unsqueeze(-1) age = age.unsqueeze(0) distance_age = (age - age_t).abs() distance = torch.zeros_like(distance_age) if 'age' in self.demo_keys: distance += distance_age if 'bmi' in self.demo_keys: # BMI adjustment bmi = (bmi - mu_bmi) / std_bmi bmi_t = bmi.unsqueeze(-1) bmi = bmi.unsqueeze(0) distance_bmi = (bmi - bmi_t).abs() distance += distance_bmi if 'age' in self.demo_keys and 'bmi' in self.demo_keys: distance = distance / 2 if 'sex' in self.demo_keys: diff_sex_matrix = 1.0 - (sex.unsqueeze(0) == sex.unsqueeze(1)).float() # Same sex: 0, different sex: 1 distance += diff_sex_matrix distance += torch.eye(N, device=age.device) * 1e8 mask_seg = mask.view(N, -1, self.cycle_len).any(-1) # (N, num_seg) distance = distance.unsqueeze(0).repeat(num_cycle,1,1) # (num_cycle, N, N) Same for all cycles mask1 = mask_seg.permute(1,0) # (num_cycle, N) label = F.softmax((- distance / self.upsilon), dim=-1) # (num_cycle, N, N) if mask is not None: xf = xf.masked_fill(~mask.unsqueeze(-1), 0.0) min_length = mask.sum(dim=1).min() num_seg = min_length // self.cycle_len # number of cycles that can be fully formed for all sequences in the batch if self.use_transformer: cycle_mode = 'mean' else: cycle_mode = 'last' seq1 = self._get_cycle_values(xf, mask, mode=cycle_mode) # (N, num_seg + 1, D) # if self.use_transformer: # seq1 = xf.view(N, -1, self.cycle_len, D).mean(2) # (N, num_seg, D) # seq1 = self.contr_head(seq1) # seq1 = F.normalize(seq1, p=2, dim=-1) # # else: # # Last feature # if not min_length == L: # mask_len = mask.sum(dim=1) # seq1_last = xf[torch.arange(N, device=xf.device), mask_len - 1] # seq1_last = self.contr_head(seq1_last) # seq1_last = F.normalize(seq1_last, p=2, dim=-1) # if num_seg < 1: # seq1 = seq1_last.unsqueeze(1) # else: # min_length = num_seg * self.cycle_len # seq_trunc = xf[:, :min_length, :] # seq1 = seq_trunc[:,(self.cycle_len-1)::self.cycle_len] # B_temp, S_temp, D_temp = seq1.shape # seq1 = seq1.reshape(-1, D_temp) # flatten batch and segments # seq1 = self.contr_head(seq1) # seq1 = seq1.reshape(B_temp, S_temp, -1) # reshape back # seq1 = F.normalize(seq1, p=2, dim=-1) # (N, num_seg, D) # # Last feature # if not min_length == L: # seq1 = torch.cat([seq1, seq1_last.unsqueeze(1)], dim=1) # (N, num_seg+1, D) koleo_loss = 0.0 for i in range(seq1.shape[1]): koleo_loss += self.koleo_loss(seq1[i]) koleo_loss = koleo_loss / seq1.shape[1] # Patient by patient similarity seq1 = seq1.permute(1,0,2).contiguous() # (num_seg + 1, N, D) sims_pos = seq1 @ seq1.permute(0,2,1) - torch.eye(N, device=age.device).unsqueeze(0) * 1e8 # (num_seg + 1, N, N) sims = sims_pos / self.tau # (num_seg + 1, N, N) # print("## Similarity", sims[0,0]) loss1, cnt = 0, 0 # For each segment, calculate the loss only among valid sequences # (with enough length) to avoid noise from padding for i in range(num_seg): # Skip any segments with too few valid sequences. if mask1[i].sum() <= 4: continue sims1 = sims[i][mask1[i]].masked_fill(~mask1[i], -1e8) label1 = label[i][mask1[i]].masked_fill(~mask1[i], 0.0) loss1 += nn.CrossEntropyLoss(reduction='mean')(sims1, label1) cnt += 1 # Additionally, calculate loss for the last segment. # num_seg is the number of full segments. if not self.use_transformer and not min_length == L: loss1 += nn.CrossEntropyLoss(reduction='mean')(sims[-1], label[0]) cnt += 1 loss1 = loss1 / cnt if cnt > 0 else loss1 loss1 = loss1 + self.koleo_weight * koleo_loss return loss1 # def init_teacher(self): # raise RuntimeError("init_teacher was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.") # def _forward_dino(self, sig, mask): # raise RuntimeError("_forward_dino was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.") # def _forward_simclr(self, sig, mask): # raise RuntimeError("_forward_simclr was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use SimCLR.") # def _simclr_representation(self, sig, mask): # raise RuntimeError("_simclr_representation was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use SimCLR.") def _get_cycle_values(self, xf, mask, mode='last'): """ Get the last/mean features of each cycle. "mode" is either 'last' or 'mean'. 'last' is used in the paper (Mamba), but 'mean' can be an alternative (Transformer) """ if mask is not None: xf = xf.masked_fill(~mask.unsqueeze(-1), 0.0) min_length = mask.sum(dim=1).min() num_seg = min_length // self.cycle_len N, L, D = xf.shape if mode == 'mean': seq1 = xf.view(N, -1, self.cycle_len, D).mean(2) # (N, num_seg, D) seq1 = self.contr_head(seq1) seq1 = F.normalize(seq1, p=2, dim=-1) return seq1 # Last feature if not min_length == L: mask_len = mask.sum(dim=1) seq1_last = xf[torch.arange(N, device=xf.device), mask_len - 1] seq1_last = self.contr_head(seq1_last) seq1_last = F.normalize(seq1_last, p=2, dim=-1) if num_seg < 1: seq1 = seq1_last.unsqueeze(1) else: min_length = num_seg * self.cycle_len seq_trunc = xf[:, :min_length, :] seq1 = seq_trunc[:,(self.cycle_len-1)::self.cycle_len] B_temp, S_temp, D_temp = seq1.shape seq1 = seq1.reshape(-1, D_temp) # flatten batch and segments seq1 = self.contr_head(seq1) seq1 = seq1.reshape(B_temp, S_temp, -1) # reshape back seq1 = F.normalize(seq1, p=2, dim=-1) # (N, num_seg, D) # Last feature if not min_length == L: seq1 = torch.cat([seq1, seq1_last.unsqueeze(1)], dim=1) # (N, num_seg+1, D) return seq1 # def _masked_mean(self, x, mask): # if mask is None: # return x.mean(1) # mask = mask.to(dtype=x.dtype) # denom = mask.sum(dim=1, keepdim=True).clamp(min=1.0) # return (x * mask.unsqueeze(-1)).sum(dim=1) / denom # def _dino_representation(self, sig, mask, use_teacher=False): # raise RuntimeError("_dino_representation was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.") # def _forward_features_with_modules(self, inputs, mask, encoder, fuser, proj, sequence_encoder): # raise RuntimeError("_forward_features_with_modules was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO teacher forward.") # def dino_loss(self, student_out, teacher_out, fwd=False): # raise RuntimeError("dino_loss was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.") # @torch.no_grad() # def _update_dino_center(self, teacher_out_f, teacher_out_r): # raise RuntimeError("_update_dino_center was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.") # @torch.no_grad() # def update_teacher(self, momentum=None): # raise RuntimeError("update_teacher was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.") # def _freeze_teacher(self): # raise RuntimeError("_freeze_teacher was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.") # def simclr_loss(self, z1, z2): # raise RuntimeError("simclr_loss was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use SimCLR.") def koleo_loss(self, x, ys=None, eps=1e-8): """ Compute KoLeo (Kozachenko-Leonenko) regularization loss. Encourages uniform distribution of representations on the hypersphere. Args: x: Tensor of shape (N, D) where N is batch size and D is feature dimension eps: Small constant for numerical stability Returns: KoLeo loss (scalar) """ # Normalize representations to unit hypersphere x = F.normalize(x, p=2, dim=-1) # Compute pairwise distances # Using efficient computation: ||a - b||^2 = ||a||^2 + ||b||^2 - 2 # Since normalized, ||a||^2 = ||b||^2 = 1, so ||a - b||^2 = 2 - 2 gram_matrix = torch.matmul(x, x.t()) # (N, N) sq_distances = 2 - 2 * gram_matrix # (N, N) sq_distances = torch.clamp(sq_distances, min=0) # numerical stability # Set diagonal to infinity (exclude self-distance) sq_distances = sq_distances + torch.eye(x.shape[0], device=x.device) * 1e10 if ys is not None: for y in ys: y = F.normalize(y, p=2, dim=-1) y_distance = (x - y).pow(2).sum(dim=-1, keepdim=True) # (N, 1) sq_distances = torch.cat([sq_distances, y_distance], dim=1) # (N, N + len(ys)) # Find minimum distance for each sample (nearest neighbor) min_sq_distances, _ = torch.min(sq_distances, dim=1) # (N,) # KoLeo loss is negative log of minimum distances # We want to maximize minimum distances -> minimize negative log distances = torch.sqrt(min_sq_distances + eps) koleo_loss = -torch.log(distances + eps).mean() return koleo_loss class SleepEventLSTMClassifier(nn.Module): def __init__(self, input_dim=384*6+512*2, embed_dim=128, num_heads=4, num_layers=1, num_classes=5, pooling_head=4, dropout=0.3, max_seq_length=128): super().__init__() # Define spatial pooling # self.spatial_pooling = AttentionPooling(embed_dim, num_heads=pooling_head, dropout=dropout) # Set max sequence length # if max_seq_length is None: # max_seq_length = 20000 # self.positional_encoding = PositionalEncoding(max_seq_length, embed_dim) self.layer_norm = nn.LayerNorm(input_dim) self.proj = nn.Linear(input_dim, embed_dim) # Transformer encoder for spatial modeling 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) # LSTM for temporal modeling lstm_dropout = dropout if num_layers > 1 else 0.0 self.lstm = nn.LSTM(input_size=embed_dim, hidden_size=embed_dim//2, num_layers=num_layers, batch_first=True, dropout=lstm_dropout, bidirectional=True) # Fully connected layer for classification self.fc = nn.Linear(embed_dim, num_classes) def forward(self, x, mask, demo=None): # B, C, S, E = x.shape # Rearrange for spatial pooling # x = rearrange(x, 'b c s e -> (b s) c e') # # Prepare the mask for spatial pooling # mask_spatial = mask[:, :, 0] # mask_spatial = mask_spatial.unsqueeze(1).expand(-1, S, -1) # mask_spatial = rearrange(mask_spatial, 'b t c -> (b t) c') # # Ensure the mask is boolean # if mask_spatial.dtype != torch.bool: # mask_spatial = mask_spatial.to(dtype=torch.bool) # # Apply spatial pooling # x = self.spatial_pooling(x, mask_spatial) # Reshape to (B, S, E) after pooling # x = x.view(B, S, E) # Apply positional encoding and layer normalization # x = self.positional_encoding(x) x = self.layer_norm(x) x = self.proj(x) # Apply transformer encoder for spatial modeling # mask_temporal = mask[:, 0, :] mask_temporal = ~mask # 1 should be padding, 0 should be valid values x = self.transformer_encoder(x, src_key_padding_mask=mask_temporal) # Apply LSTM for temporal modeling x, _ = self.lstm(x) # Shape: (B, S, E) # Apply the final fully connected layer for classification x = self.fc(x) # Shape: (B, S, num_classes) return x # Return mask along temporal dimension # From SleepFM-clinical import torch.nn.utils.rnn as rnn_utils class DiagnosisFinetuneFullLSTMCOXPHWithDemo(nn.Module): def __init__(self, input_dim, embed_dim=128, num_layers=2, num_classes=6, dropout=0.1): super(DiagnosisFinetuneFullLSTMCOXPHWithDemo, self).__init__() self.layer_norm = nn.LayerNorm(input_dim) self.proj = nn.Linear(input_dim, embed_dim) self.lstm = nn.LSTM( input_size=embed_dim, hidden_size=embed_dim // 2, num_layers=num_layers, batch_first=True, bidirectional=True ) self.demo_embedding = nn.Sequential( nn.Linear(2, embed_dim // 4), nn.ReLU(), nn.Dropout(dropout) ) # Update the input size of the disease head to account for embedded demo features self.head = nn.Linear(embed_dim + embed_dim // 4, num_classes) # embed_dim // 4 from demo features def forward(self, x, mask, demo): # B, C, S, E = x.shape # C: Modalities # # Rearrange for spatial pooling # x = rearrange(x, 'b c s e -> (b s) c e') # mask_spatial = mask[:, :, 0] # mask_spatial = mask_spatial.unsqueeze(1).expand(-1, S, -1) # mask_spatial = rearrange(mask_spatial, 'b t c -> (b t) c') # # Ensure the mask is boolean # mask_spatial = mask_spatial.to(dtype=torch.bool) # # Apply spatial pooling # x = self.spatial_pooling(x, mask_spatial) # x = x.view(B, S, E) # # Rearrange mask for temporal pooling # # mask_temporal = mask[:, 0, :] # Use the mask along the temporal dimension # # Compute the lengths for packing based on the mask # lengths = (mask_temporal == 0).sum(dim=1).cpu() x = self.layer_norm(x) x = self.proj(x) # lengths = (mask == 1).sum(dim=1).cpu() # (B,) lengths = mask.sum(dim=1).long().cpu() # (B,) B = x.shape[0] # Pack the sequence for LSTM packed_x = rnn_utils.pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False) # Pass the packed sequence through the LSTM packed_out, _ = self.lstm(packed_x) # Unpack the sequence x, _ = rnn_utils.pad_packed_sequence(packed_out, batch_first=True) # Apply mean pooling only over valid lengths x = torch.stack([x[i, :lengths[i]].mean(dim=0) for i in range(B)]) age = (demo['age'] - 63.14)/11.22 age = torch.clamp(age, -3.0, 3.0) sex = demo['sex'] - 1 # 0 or 1 demo_features = torch.cat([age.unsqueeze(1), sex.unsqueeze(1)], dim=1) demo_features = demo_features.to(x.device, dtype=x.dtype) # Embed the demo features demo_embed = self.demo_embedding(demo_features) # (B, embed_dim // 4) # Concatenate the demo features with the LSTM output x = torch.cat([x, demo_embed], dim=1) # (B, embed_dim + embed_dim // 4) # Apply the disease head to get hazards hazards = self.head(x) return hazards class DiagnosisFinetuneDemoOnly(nn.Module): def __init__(self, embed_dim=128, num_layers=2, num_classes=6, dropout=0.1): super().__init__() self.demo_embedding = nn.Sequential( nn.Linear(2, embed_dim // 4), nn.ReLU(), nn.Dropout(dropout) ) # Update the input size of the disease head to account for embedded demo features self.head = nn.Linear(embed_dim // 4, num_classes) # embed_dim // 4 from demo features def forward(self, x, mask, demo): age = (demo['age'] - 63.14)/11.22 age = torch.clamp(age, -3.0, 3.0) sex = demo['sex'] - 1 # 0 or 1 demo_features = torch.cat([age.unsqueeze(1), sex.unsqueeze(1)], dim=1) demo_features = demo_features.to(x.device, dtype=x.dtype) # Embed the demo features demo_embed = self.demo_embedding(demo_features) # (B, embed_dim // 4) # Apply the disease head to get hazards hazards = self.head(demo_embed) return hazards