Keondopark

SDB segmentation update to report Class-F1

6196ce4 3 days ago

20.3 kB

from functools import partial
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import random

from models.encoders import MultiEncoder
from models.decoders import Decoder
from models.fusion import FusionUpsampleBlock
from models.transformer_modules import RevIN
from configs.global_cfg import TARGET_RESAMPLE



class MultiMaskedAutoencoderConv1d(nn.Module):
    """ Masked Autoencoder with VisionTransformer backbone
    """
    def __init__(self, sig_len=6000, patch_size=10,
                 embed_dim=[32,64,128,256], depth=[2,2,4], specific_depth=[2,2,2], num_heads=8,
                 decoder_embed_dim=256, decoder_depth=6, decoder_num_heads=8,
                 mlp_ratio=4, norm_layer=nn.LayerNorm,
                 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,
                 mask_ratio=0.5,
                 contrastive_training=False, koleo_training=False, pos_enc='sinusoidal'):

        super().__init__()         
        self.patch_size = patch_size        
        self.modalities = modalities        
        self.num_experts = num_experts
        self.num_experts_per = num_experts_per        
        self.embed_dim = embed_dim
        self.depth = depth
        self.specific_depth = specific_depth
        self.mask_ratio = mask_ratio
        self.encoder = MultiEncoder(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,
                 mask_ratio=mask_ratio, pos_enc=pos_enc, use_shared=use_shared, use_moe=use_moe,)

        self.num_patches = self.encoder.num_patches                
        
        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.use_shared = use_shared
        self.use_moe = use_moe
        
        self.contrastive_training = contrastive_training
        self.koleo_training = koleo_training

        # contrastive loss hyperparameters (InfoNCE style)
        self.contrastive_temp = 0.07
        
        self.fusion = nn.ModuleDict()                
        self.decoders = nn.ModuleDict()
        self.upsampler = nn.ModuleDict()
        

        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)
        
        for m in self.modalities:
            decoder_pred_dim = self.patch_size if m != 'SaO2' else 1
            self.decoders[m] = Decoder(num_patches=self.num_patches, encoder_embed_dim=embed_dim[2], decoder_embed_dim=decoder_embed_dim, 
                    decoder_depth=decoder_depth, decoder_num_heads=decoder_num_heads,
                    mlp_ratio=mlp_ratio, norm_layer=norm_layer, decoder_pred_dim=decoder_pred_dim, pos_enc=pos_enc)                        
        
        
        
        self.apply(self._init_weights)

    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
            ys: Optional list of tensors to also consider for distance computation
            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<a,b>
        # Since normalized, ||a||^2 = ||b||^2 = 1, so ||a - b||^2 = 2 - 2<a,b>
        gram_matrix = torch.matmul(x, x.t())  # (N, N)
        sq_distances = 2 - 2 * gram_matrix  # (N, N)
        sq_distances = torch.clamp(sq_distances, min=eps, max=4.0)  # numerical stability, max distance on unit sphere is 2
        
        # Set diagonal to large value (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)
                y_distance = torch.clamp(y_distance, min=eps, max=4.0)
                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)
        
        # Additional clamping on distances before log
        distances = torch.clamp(distances, min=eps)
        
        koleo_loss = -torch.log(distances).mean()
        
        return koleo_loss

    def forward(self, inputs, mask_dict=None, ids_keep_dict=None):
        # embed patches
        # m1, m2 = self.modalities[0], self.modalities[1]
        modalities_to_use = list(inputs.keys())
        # random.shuffle(modalities_to_use)
        
        m1, m2 = modalities_to_use[0], modalities_to_use[1]
        if mask_dict is None or ids_keep_dict is None:
            mask_dict, ids_keep_dict = self.random_masking(inputs[m1], self.mask_ratio)       
        
        # ids_shuffle, mask, ids_keep, ids_restore = self.random_masking(inputs[m1], mask_ratio)
        
        N = inputs[m1].shape[0]
        d = self.embed_dim[2]
        inputs_ = dict()
        if self.use_revin:
            for i, m in enumerate(modalities_to_use):
                # mask_m = mask_dict[m].reshape(N, -1, 1)                
                # mask_raw_m = torch.tile(mask_m, (1,  1, self.patch_size)).reshape(N, -1, 1)
                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], None, mode='norm')
        else:
            inputs_ = inputs

        preds = dict()
        latents_shared, latents_specific, gating_scores = self.encoder(inputs_, mask_dict, ids_keep_dict, modalities_to_use)
        
        latents_fused, upsampled = self.fuser(latents_shared, latents_specific)
        
        if self.mask_ratio == 0:
            losses = dict()
            for m in modalities_to_use:                                     
                losses[m] = 0.0                          
            
        else:
            for i, m in enumerate(modalities_to_use):                     
                ids_keep_m = ids_keep_dict[m] # (N, L * large portion)        
                preds[m] = self.decoders[m](upsampled[i], ids_keep_m)     
            # Compute reconstruction losses
            losses = self.forward_loss(inputs_, preds, mask_dict, ids_keep_dict, modalities_to_use)

        # Contrastive learning between latents_shared and latents_specific
        # Positives: (latents_shared[m1], latents_shared[m2]) for the same timestep (index)
        # Negatives: (latents_shared[m1], latents_specific[m1]) for the same timestep
        contrastive_loss = 0.0
        koleo_loss = 0.0
        num_contrast_pairs = 0
        
        ctrst_grp = 30 # (seconds)
        ctrst_grp = (ctrst_grp * (100 / self.patch_size)) * (1 - self.mask_ratio) # 0.5 is mask ratio
        # How many patches form a patch group for contrastive learning?
        ctrst_grp = int(ctrst_grp // (5 * 2 ** (len(self.depth) - 2))) # The first downsampling is 5x, then 2x for each level
    
        
        if self.contrastive_training:      

            z_shared, z_specific = [], []
            for i, m in enumerate(modalities_to_use):
                shared_m = latents_shared[-1][i]
                N, L, D = shared_m.shape
                Lm = shared_m.reshape(-1, ctrst_grp*D)
                # Lm = self.pooler1[m](Lm)
                Lm = F.normalize(Lm, dim=-1)                
                z_shared.append(Lm)  # (N, L, D')

                specific_m = latents_specific[-1][i]
                # N, L, D = specific_m.shape
                # specific_m = specific_m.reshape(N, L, 4, D).mean(dim=2)
                Lm = specific_m.reshape(-1, ctrst_grp*D)
                # Lm = self.pooler2[m](Lm)                
                Lm = F.normalize(Lm, dim=-1)                
                z_specific.append(Lm)  # (N, L, D')

            
            for i, m1 in enumerate(modalities_to_use):                  
                shared_specific = []  
                z_anchor = z_shared[i]
                
                
                sims_neg1 = z_anchor @ z_anchor.permute(1,0) # (N*L, N*L)                
                diag_idx = torch.arange(sims_neg1.size(1), device=sims_neg1.device)
                sims_neg1[diag_idx, diag_idx] = float('-inf')
                
                
                z_pos = [z for k, z in enumerate(z_shared) if k != i]
                z_pos = torch.mean(torch.stack(z_pos, dim=0), dim=0)  # Average over other modalities
                
                sims_pos = torch.sum(z_anchor * z_pos, dim=-1, keepdim=True)  # (N*L, 1)            
                
                if self.use_shared:
                    zs = z_specific[i]
                    sims_neg2 = torch.sum(z_anchor * zs, dim=-1, keepdim=True)
                    shared_specific.append(sims_neg2)
                    sims_neg2 = torch.cat(shared_specific, dim=-1) # (N*L, num_modalities)              

                    logits = torch.cat([sims_pos, sims_neg1, sims_neg2], dim=-1) / (self.contrastive_temp + 1e-8)                     
                else:
                    logits = torch.cat([sims_pos, sims_neg1], dim=-1) / (self.contrastive_temp + 1e-8)
                # logits = torch.cat([sims_pos, sims_neg1, sims_neg2], dim=-1) * logit_scale 
                labels = torch.zeros(logits.shape[0], dtype=torch.long, device=logits.device)                    
                loss_c = F.cross_entropy(logits, labels)
                contrastive_loss = contrastive_loss + loss_c
                num_contrast_pairs += 1
            

            if num_contrast_pairs > 0:
                contrastive_loss = contrastive_loss / float(num_contrast_pairs)
        
        if self.koleo_training: 
            # Compute KoLeo regularization on shared representations
            # This encourages uniform distribution of shared features on the hypersphere
            # for i in range(len(modalities_to_use)):
            #     shared_flat = z_shared[i]
                          
               
            #     specific_flats = [z_specific[i]]
            #     # combined_flat = torch.cat([shared_flat, specific_m1_flat, specific_m2_flat], dim=0)
                        
            #     koleo_loss = koleo_loss + self.koleo_loss(shared_flat, ys=specific_flats)
            koleo_loss = self.koleo_loss(torch.cat(z_shared, dim=0), ys=torch.cat(z_specific, dim=0))
         

        
        aux_loss = self.load_balancing_loss(gating_scores)
        # return losses, aux_loss_lg, aux_loss_sm, preds, mask_dict, ids_keep_dict
        return losses, aux_loss, contrastive_loss, koleo_loss, preds, mask_dict, ids_keep_dict

    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 patchify(self, sig):
        """
        sig: (N, L, 1)
        patchfied: (N, L, patch_size)
        """
        # p = self.patch_size
        B, L, _ = sig.shape

        num_patches = L // self.patch_size

        x = sig.reshape(shape=(B, num_patches, self.patch_size))

        return x
    
    def random_masking(self, x, mask_ratio=0.5):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """

        mask_dict, ids_keep_dict = dict(), dict()
        
        # If no mask is used. Keep all patches.
        if mask_ratio == 0:
            for m in self.modalities:
                mask_dict[m] = torch.zeros(x.shape[0], x.shape[1] // self.patch_size, device=x.device, dtype=x.dtype)
                ids_keep_dict[m] = torch.arange(x.shape[1] // self.patch_size, device=x.device).unsqueeze(0).repeat([x.shape[0], 1])
            return mask_dict, ids_keep_dict

        N = x.shape[0]        
        L = x.shape[1] // self.patch_size # number of large patches        

        
        len_keep = int(L * (1 - mask_ratio))  # Number of keep large patches

        # noise = torch.rand(N, L_lg, device=x.device)  # noise in [0, 1]
        # ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove     
        
        ids_ordered = torch.arange(L, device=x.device).unsqueeze(0).repeat([N, 1])
        # ids_ordered = ids_ordered.view([N, L //2, 2])
        rand_ids = torch.randint(0, 2, [N, len_keep, 1], device=x.device)
        rand_ids2 = torch.randint(0, 2, [N, len_keep, 1], device=x.device)
        ids_keep_m1 = torch.gather(ids_ordered.view([N, len_keep, -1]), index=rand_ids, dim=-1).view([N, -1])
        ids_keep_m2 = torch.gather(ids_ordered.view([N, len_keep, -1]), index=rand_ids2, dim=-1).view([N, -1])
        
        num_patches_slow = (200 // self.patch_size)
        len_keep_slow = int((L // num_patches_slow) * (1 - mask_ratio))  # Number of keep large patches for slow modalities
        
        rand_ids = torch.randint(0, 2, [N, len_keep_slow, 1], device=x.device)
        ids_keep_slow = torch.gather(ids_ordered.view([N, len_keep_slow, -1, num_patches_slow]), 
                                index=rand_ids.unsqueeze(-1).repeat([1,1,1,num_patches_slow]), 
                                dim=2).view([N, -1])
        
        
        # ids_keep_m1 = ids[:, 0::2]  ## Large keep area, modality 1
        # ids_keep_m2 = ids[:, 1::2]  ## Large keep area, modality 2
        
         # Mask construction for large modeling area
        mask_m1 = torch.ones([N, L], device=x.device, dtype=x.dtype)
        src = torch.zeros([N, len_keep], device=x.device, dtype=x.dtype)
        mask_m1.scatter_(dim=1, index=ids_keep_m1, src=src)  # 0 is keep, 1 is remove, Large mask + small modeling area masked
 
        mask_m2 = torch.ones([N, L], device=x.device, dtype=x.dtype)
        src = torch.zeros([N, L - len_keep], device=x.device, dtype=x.dtype)
        mask_m2.scatter_(dim=1, index=ids_keep_m2, src=src)  # 0 is keep, 1 is remove, Large mask + small modeling area masked
        
        mask_slow = torch.ones([N, L], device=x.device, dtype=x.dtype)
        src = torch.zeros([N, L - len_keep], device=x.device, dtype=x.dtype)
        mask_slow.scatter_(dim=1, index=ids_keep_slow, src=src)  # 0 is keep, 1 is remove, Large mask + small modeling area masked
        # print(rand_ids)
        # print(ids_keep_slow)
        

        for i, m in enumerate(self.modalities):
            idx = (i % 2) + 1            
            if m in ['RESP','SaO2']:
            # if False: 
                ids_keep_dict[m] = ids_keep_slow
                mask_dict[m] = mask_slow  # Both use the same mask
            else:
                ids_keep_dict[m] = ids_keep_m1 if idx == 1 else ids_keep_m2            
                mask_dict[m] = mask_m1 if idx == 1 else mask_m2            


        return mask_dict, ids_keep_dict

    def forward_loss(self, inputs, preds, mask_dict, ids_keep_dict, modalities_to_use):
        """
        inputs: m * [N, L]
        pred: [N, L, p*p*3]
        mask: [N, L], 0 is keep, 1 is remove, 
        """
        losses = dict()        
        # lg_ids, sm_ids = ids_keep_dict['ids_lg'], ids_keep_dict['ids_sm']
        # for m in inputs.keys():
        for i, m in enumerate(modalities_to_use):            
            
            sig, pred, mask= inputs[m], preds[m], mask_dict[m]
            
            N, L, _ = sig.shape

            if m == 'SaO2':
                repeat_factor = 100 // self.patch_size
                target = sig.repeat_interleave(repeat_factor, dim=1)
                pred = pred.view(target.shape)            
                
                loss = (pred - target) ** 2               
                loss = (loss.mean(dim=-1) * mask).sum() / mask.sum()            

            else:            
                target = self.patchify(sig) # (N, L, patch_size)            
                pred = pred.view(target.shape)            
                
                sig_smooth = F.avg_pool1d(sig.view(N, -1, L), 
                                kernel_size=11, 
                                stride=1, 
                                padding=5, 
                                count_include_pad=False).view(N, L, -1)
                
                target_smooth = self.patchify(sig_smooth)
                loss_smooth = (pred - target_smooth) ** 2               
                loss_smooth = (loss_smooth.mean(dim=-1) * mask).sum() / mask.sum()            
                
                loss = loss_smooth
            
            
            losses[m] = loss
        return losses
    
    def load_balancing_loss(self, gating_scores):
        """
        Compute the load balancing loss for the gating scores.
        ref: https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py#L591
        """
        num_moe = len(gating_scores)
        if num_moe > 0:
            overall_loss = 0.0
            for i in range(num_moe):
                gating_scores_i = gating_scores[i]  # stop gradient
                # gating_scores = torch.stack(gating_scores, dim=0)  # [num_MoEblocks, B, num_experts]
                routing_weights = gating_scores_i.softmax(dim=-1)  # [B, L, num_experts]
                _, selected_experts = torch.topk(routing_weights, self.num_experts_per, dim=-1) # Index of selected experts
                expert_mask = torch.nn.functional.one_hot(selected_experts, self.num_experts) # [B, L, num_experts_per, num_experts]

                # Compute the average number of data points routed to each experts
                # tokens_per_expert = torch.mean(expert_mask.float(), dim=1) # [B, num_experts]

                # # Compute the average probability of routing to these experts
                # router_prob_per_expert = torch.mean(routing_weights, dim=0) # [num_experts]

                # overall_loss += torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(1))
                tokens_per_expert = torch.mean(expert_mask.float(), dim=(0,1,2)) # [num_experts]
                router_prob_per_expert = torch.mean(routing_weights, dim=(0,1)) # [num_experts]
                overall_loss += torch.mean(tokens_per_expert * router_prob_per_expert)
            
            return overall_loss / num_moe
        else:
            return 0.0

    

	from functools import partial
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	import numpy as np
	import random

	from models.encoders import MultiEncoder
	from models.decoders import Decoder
	from models.fusion import FusionUpsampleBlock
	from models.transformer_modules import RevIN
	from configs.global_cfg import TARGET_RESAMPLE



	class MultiMaskedAutoencoderConv1d(nn.Module):
	""" Masked Autoencoder with VisionTransformer backbone
	"""
	def __init__(self, sig_len=6000, patch_size=10,
	embed_dim=[32,64,128,256], depth=[2,2,4], specific_depth=[2,2,2], num_heads=8,
	decoder_embed_dim=256, decoder_depth=6, decoder_num_heads=8,
	mlp_ratio=4, norm_layer=nn.LayerNorm,
	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,
	mask_ratio=0.5,
	contrastive_training=False, koleo_training=False, pos_enc='sinusoidal'):

	super().__init__()
	self.patch_size = patch_size
	self.modalities = modalities
	self.num_experts = num_experts
	self.num_experts_per = num_experts_per
	self.embed_dim = embed_dim
	self.depth = depth
	self.specific_depth = specific_depth
	self.mask_ratio = mask_ratio
	self.encoder = MultiEncoder(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,
	mask_ratio=mask_ratio, pos_enc=pos_enc, use_shared=use_shared, use_moe=use_moe,)

	self.num_patches = self.encoder.num_patches

	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.use_shared = use_shared
	self.use_moe = use_moe

	self.contrastive_training = contrastive_training
	self.koleo_training = koleo_training

	# contrastive loss hyperparameters (InfoNCE style)
	self.contrastive_temp = 0.07

	self.fusion = nn.ModuleDict()
	self.decoders = nn.ModuleDict()
	self.upsampler = nn.ModuleDict()


	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)

	for m in self.modalities:
	decoder_pred_dim = self.patch_size if m != 'SaO2' else 1
	self.decoders[m] = Decoder(num_patches=self.num_patches, encoder_embed_dim=embed_dim[2], decoder_embed_dim=decoder_embed_dim,
	decoder_depth=decoder_depth, decoder_num_heads=decoder_num_heads,
	mlp_ratio=mlp_ratio, norm_layer=norm_layer, decoder_pred_dim=decoder_pred_dim, pos_enc=pos_enc)



	self.apply(self._init_weights)

	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
	ys: Optional list of tensors to also consider for distance computation
	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<a,b>
	# Since normalized, \|\|a\|\|^2 = \|\|b\|\|^2 = 1, so \|\|a - b\|\|^2 = 2 - 2<a,b>
	gram_matrix = torch.matmul(x, x.t()) # (N, N)
	sq_distances = 2 - 2 * gram_matrix # (N, N)
	sq_distances = torch.clamp(sq_distances, min=eps, max=4.0) # numerical stability, max distance on unit sphere is 2

	# Set diagonal to large value (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)
	y_distance = torch.clamp(y_distance, min=eps, max=4.0)
	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)

	# Additional clamping on distances before log
	distances = torch.clamp(distances, min=eps)

	koleo_loss = -torch.log(distances).mean()

	return koleo_loss

	def forward(self, inputs, mask_dict=None, ids_keep_dict=None):
	# embed patches
	# m1, m2 = self.modalities[0], self.modalities[1]
	modalities_to_use = list(inputs.keys())
	# random.shuffle(modalities_to_use)

	m1, m2 = modalities_to_use[0], modalities_to_use[1]
	if mask_dict is None or ids_keep_dict is None:
	mask_dict, ids_keep_dict = self.random_masking(inputs[m1], self.mask_ratio)

	# ids_shuffle, mask, ids_keep, ids_restore = self.random_masking(inputs[m1], mask_ratio)

	N = inputs[m1].shape[0]
	d = self.embed_dim[2]
	inputs_ = dict()
	if self.use_revin:
	for i, m in enumerate(modalities_to_use):
	# mask_m = mask_dict[m].reshape(N, -1, 1)
	# mask_raw_m = torch.tile(mask_m, (1, 1, self.patch_size)).reshape(N, -1, 1)
	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], None, mode='norm')
	else:
	inputs_ = inputs

	preds = dict()
	latents_shared, latents_specific, gating_scores = self.encoder(inputs_, mask_dict, ids_keep_dict, modalities_to_use)

	latents_fused, upsampled = self.fuser(latents_shared, latents_specific)

	if self.mask_ratio == 0:
	losses = dict()
	for m in modalities_to_use:
	losses[m] = 0.0

	else:
	for i, m in enumerate(modalities_to_use):
	ids_keep_m = ids_keep_dict[m] # (N, L * large portion)
	preds[m] = self.decoders[m](upsampled[i], ids_keep_m)
	# Compute reconstruction losses
	losses = self.forward_loss(inputs_, preds, mask_dict, ids_keep_dict, modalities_to_use)

	# Contrastive learning between latents_shared and latents_specific
	# Positives: (latents_shared[m1], latents_shared[m2]) for the same timestep (index)
	# Negatives: (latents_shared[m1], latents_specific[m1]) for the same timestep
	contrastive_loss = 0.0
	koleo_loss = 0.0
	num_contrast_pairs = 0

	ctrst_grp = 30 # (seconds)
	ctrst_grp = (ctrst_grp * (100 / self.patch_size)) * (1 - self.mask_ratio) # 0.5 is mask ratio
	# How many patches form a patch group for contrastive learning?
	ctrst_grp = int(ctrst_grp // (5 * 2 ** (len(self.depth) - 2))) # The first downsampling is 5x, then 2x for each level


	if self.contrastive_training:

	z_shared, z_specific = [], []
	for i, m in enumerate(modalities_to_use):
	shared_m = latents_shared[-1][i]
	N, L, D = shared_m.shape
	Lm = shared_m.reshape(-1, ctrst_grp*D)
	# Lm = self.pooler1[m](Lm)
	Lm = F.normalize(Lm, dim=-1)
	z_shared.append(Lm) # (N, L, D')

	specific_m = latents_specific[-1][i]
	# N, L, D = specific_m.shape
	# specific_m = specific_m.reshape(N, L, 4, D).mean(dim=2)
	Lm = specific_m.reshape(-1, ctrst_grp*D)
	# Lm = self.pooler2[m](Lm)
	Lm = F.normalize(Lm, dim=-1)
	z_specific.append(Lm) # (N, L, D')


	for i, m1 in enumerate(modalities_to_use):
	shared_specific = []
	z_anchor = z_shared[i]


	sims_neg1 = z_anchor @ z_anchor.permute(1,0) # (NL, NL)
	diag_idx = torch.arange(sims_neg1.size(1), device=sims_neg1.device)
	sims_neg1[diag_idx, diag_idx] = float('-inf')


	z_pos = [z for k, z in enumerate(z_shared) if k != i]
	z_pos = torch.mean(torch.stack(z_pos, dim=0), dim=0) # Average over other modalities

	sims_pos = torch.sum(z_anchor * z_pos, dim=-1, keepdim=True) # (N*L, 1)

	if self.use_shared:
	zs = z_specific[i]
	sims_neg2 = torch.sum(z_anchor * zs, dim=-1, keepdim=True)
	shared_specific.append(sims_neg2)
	sims_neg2 = torch.cat(shared_specific, dim=-1) # (N*L, num_modalities)

	logits = torch.cat([sims_pos, sims_neg1, sims_neg2], dim=-1) / (self.contrastive_temp + 1e-8)
	else:
	logits = torch.cat([sims_pos, sims_neg1], dim=-1) / (self.contrastive_temp + 1e-8)
	# logits = torch.cat([sims_pos, sims_neg1, sims_neg2], dim=-1) * logit_scale
	labels = torch.zeros(logits.shape[0], dtype=torch.long, device=logits.device)
	loss_c = F.cross_entropy(logits, labels)
	contrastive_loss = contrastive_loss + loss_c
	num_contrast_pairs += 1


	if num_contrast_pairs > 0:
	contrastive_loss = contrastive_loss / float(num_contrast_pairs)

	if self.koleo_training:
	# Compute KoLeo regularization on shared representations
	# This encourages uniform distribution of shared features on the hypersphere
	# for i in range(len(modalities_to_use)):
	# shared_flat = z_shared[i]


	# specific_flats = [z_specific[i]]
	# # combined_flat = torch.cat([shared_flat, specific_m1_flat, specific_m2_flat], dim=0)

	# koleo_loss = koleo_loss + self.koleo_loss(shared_flat, ys=specific_flats)
	koleo_loss = self.koleo_loss(torch.cat(z_shared, dim=0), ys=torch.cat(z_specific, dim=0))



	aux_loss = self.load_balancing_loss(gating_scores)
	# return losses, aux_loss_lg, aux_loss_sm, preds, mask_dict, ids_keep_dict
	return losses, aux_loss, contrastive_loss, koleo_loss, preds, mask_dict, ids_keep_dict

	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 patchify(self, sig):
	"""
	sig: (N, L, 1)
	patchfied: (N, L, patch_size)
	"""
	# p = self.patch_size
	B, L, _ = sig.shape

	num_patches = L // self.patch_size

	x = sig.reshape(shape=(B, num_patches, self.patch_size))

	return x

	def random_masking(self, x, mask_ratio=0.5):
	"""
	Perform per-sample random masking by per-sample shuffling.
	Per-sample shuffling is done by argsort random noise.
	x: [N, L, D], sequence
	"""

	mask_dict, ids_keep_dict = dict(), dict()

	# If no mask is used. Keep all patches.
	if mask_ratio == 0:
	for m in self.modalities:
	mask_dict[m] = torch.zeros(x.shape[0], x.shape[1] // self.patch_size, device=x.device, dtype=x.dtype)
	ids_keep_dict[m] = torch.arange(x.shape[1] // self.patch_size, device=x.device).unsqueeze(0).repeat([x.shape[0], 1])
	return mask_dict, ids_keep_dict

	N = x.shape[0]
	L = x.shape[1] // self.patch_size # number of large patches


	len_keep = int(L * (1 - mask_ratio)) # Number of keep large patches

	# noise = torch.rand(N, L_lg, device=x.device) # noise in [0, 1]
	# ids_shuffle = torch.argsort(noise, dim=1) # ascend: small is keep, large is remove

	ids_ordered = torch.arange(L, device=x.device).unsqueeze(0).repeat([N, 1])
	# ids_ordered = ids_ordered.view([N, L //2, 2])
	rand_ids = torch.randint(0, 2, [N, len_keep, 1], device=x.device)
	rand_ids2 = torch.randint(0, 2, [N, len_keep, 1], device=x.device)
	ids_keep_m1 = torch.gather(ids_ordered.view([N, len_keep, -1]), index=rand_ids, dim=-1).view([N, -1])
	ids_keep_m2 = torch.gather(ids_ordered.view([N, len_keep, -1]), index=rand_ids2, dim=-1).view([N, -1])

	num_patches_slow = (200 // self.patch_size)
	len_keep_slow = int((L // num_patches_slow) * (1 - mask_ratio)) # Number of keep large patches for slow modalities

	rand_ids = torch.randint(0, 2, [N, len_keep_slow, 1], device=x.device)
	ids_keep_slow = torch.gather(ids_ordered.view([N, len_keep_slow, -1, num_patches_slow]),
	index=rand_ids.unsqueeze(-1).repeat([1,1,1,num_patches_slow]),
	dim=2).view([N, -1])


	# ids_keep_m1 = ids[:, 0::2] ## Large keep area, modality 1
	# ids_keep_m2 = ids[:, 1::2] ## Large keep area, modality 2

	# Mask construction for large modeling area
	mask_m1 = torch.ones([N, L], device=x.device, dtype=x.dtype)
	src = torch.zeros([N, len_keep], device=x.device, dtype=x.dtype)
	mask_m1.scatter_(dim=1, index=ids_keep_m1, src=src) # 0 is keep, 1 is remove, Large mask + small modeling area masked

	mask_m2 = torch.ones([N, L], device=x.device, dtype=x.dtype)
	src = torch.zeros([N, L - len_keep], device=x.device, dtype=x.dtype)
	mask_m2.scatter_(dim=1, index=ids_keep_m2, src=src) # 0 is keep, 1 is remove, Large mask + small modeling area masked

	mask_slow = torch.ones([N, L], device=x.device, dtype=x.dtype)
	src = torch.zeros([N, L - len_keep], device=x.device, dtype=x.dtype)
	mask_slow.scatter_(dim=1, index=ids_keep_slow, src=src) # 0 is keep, 1 is remove, Large mask + small modeling area masked
	# print(rand_ids)
	# print(ids_keep_slow)


	for i, m in enumerate(self.modalities):
	idx = (i % 2) + 1
	if m in ['RESP','SaO2']:
	# if False:
	ids_keep_dict[m] = ids_keep_slow
	mask_dict[m] = mask_slow # Both use the same mask
	else:
	ids_keep_dict[m] = ids_keep_m1 if idx == 1 else ids_keep_m2
	mask_dict[m] = mask_m1 if idx == 1 else mask_m2


	return mask_dict, ids_keep_dict

	def forward_loss(self, inputs, preds, mask_dict, ids_keep_dict, modalities_to_use):
	"""
	inputs: m * [N, L]
	pred: [N, L, pp3]
	mask: [N, L], 0 is keep, 1 is remove,
	"""
	losses = dict()
	# lg_ids, sm_ids = ids_keep_dict['ids_lg'], ids_keep_dict['ids_sm']
	# for m in inputs.keys():
	for i, m in enumerate(modalities_to_use):

	sig, pred, mask= inputs[m], preds[m], mask_dict[m]

	N, L, _ = sig.shape

	if m == 'SaO2':
	repeat_factor = 100 // self.patch_size
	target = sig.repeat_interleave(repeat_factor, dim=1)
	pred = pred.view(target.shape)

	loss = (pred - target) ** 2
	loss = (loss.mean(dim=-1) * mask).sum() / mask.sum()

	else:
	target = self.patchify(sig) # (N, L, patch_size)
	pred = pred.view(target.shape)

	sig_smooth = F.avg_pool1d(sig.view(N, -1, L),
	kernel_size=11,
	stride=1,
	padding=5,
	count_include_pad=False).view(N, L, -1)

	target_smooth = self.patchify(sig_smooth)
	loss_smooth = (pred - target_smooth) ** 2
	loss_smooth = (loss_smooth.mean(dim=-1) * mask).sum() / mask.sum()

	loss = loss_smooth


	losses[m] = loss
	return losses

	def load_balancing_loss(self, gating_scores):
	"""
	Compute the load balancing loss for the gating scores.
	ref: https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py#L591
	"""
	num_moe = len(gating_scores)
	if num_moe > 0:
	overall_loss = 0.0
	for i in range(num_moe):
	gating_scores_i = gating_scores[i] # stop gradient
	# gating_scores = torch.stack(gating_scores, dim=0) # [num_MoEblocks, B, num_experts]
	routing_weights = gating_scores_i.softmax(dim=-1) # [B, L, num_experts]
	_, selected_experts = torch.topk(routing_weights, self.num_experts_per, dim=-1) # Index of selected experts
	expert_mask = torch.nn.functional.one_hot(selected_experts, self.num_experts) # [B, L, num_experts_per, num_experts]

	# Compute the average number of data points routed to each experts
	# tokens_per_expert = torch.mean(expert_mask.float(), dim=1) # [B, num_experts]

	# # Compute the average probability of routing to these experts
	# router_prob_per_expert = torch.mean(routing_weights, dim=0) # [num_experts]

	# overall_loss += torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(1))
	tokens_per_expert = torch.mean(expert_mask.float(), dim=(0,1,2)) # [num_experts]
	router_prob_per_expert = torch.mean(routing_weights, dim=(0,1)) # [num_experts]
	overall_loss += torch.mean(tokens_per_expert * router_prob_per_expert)

	return overall_loss / num_moe
	else:
	return 0.0