# -*- coding: utf-8 -*- """ MDD(우울증) 딥러닝 분류기 - 모델 구조 정의 모델 제공사의 model.py 를 그대로 이식 (BHITransformer 및 하위 모듈만 사용, 학습 실험용이었던 ShallowNet 은 제외). 가중치(model/mdd/BHITransformer/models/bestModel_foldNN.pt)는 이 구조를 바꾸지 않고 그대로 로드해서 쓴다 - 이 파일 내용을 임의로 수정하면 안 된다. """ import math import torch import torch.nn as nn from timm.layers import DropPath # ---------- Positional / Time Encoding ---------- class PositionalEncoding(nn.Module): """ mode: - 'sinusoidal' : 고정식 위치(토큰 index 기반) 인코딩 - 'learnable' : 학습형 위치(토큰 index 기반) 인코딩 - 'time_fourier' : 절대 시간(초) -> Fourier features -> Linear -> D - 'time_learnable' : 절대 시간(초) -> 작은 MLP -> D (시간 범위 정규화 필요) - 'none' : 인코딩 없음 """ def __init__(self, embed_dim, mode='time_learnable', max_len=4096, time_n_freq=16, time_max_freq=8, time_min_s=0.0, time_max_s=8.0): super().__init__() self.mode = mode self.embed_dim = embed_dim self.max_len = max_len self.register_buffer('pos_embed_buf', None) self.pos_embed = None self.register_buffer('time_omega', None) self.time_proj = None self.time_mlp = None self.time_min_s = time_min_s self.time_max_s = time_max_s self.time_eps = 1e-6 if mode == 'sinusoidal': pe = torch.zeros(max_len, embed_dim) position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1) div_term = torch.exp(torch.arange(0, embed_dim, 2, dtype=torch.float32) * (-math.log(10000.0) / embed_dim)) pe[:, 0::2] = torch.sin(position * div_term) if embed_dim % 2 == 1: pe[:, 1::2] = torch.cos(position * div_term[:-1]) else: pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer('pos_embed_buf', pe.unsqueeze(0)) elif mode == 'learnable': self.pos_embed = nn.Parameter(torch.randn(1, max_len, embed_dim) * 0.02) elif mode == 'time_fourier': freqs_hz = torch.logspace(start=0, end=math.log10(time_max_freq), steps=time_n_freq) self.register_buffer('time_omega', 2 * math.pi * freqs_hz) self.time_proj = nn.Linear(2 * time_n_freq, embed_dim) elif mode == 'time_learnable': self.time_mlp = nn.Sequential( nn.Linear(1, embed_dim // 2), nn.ReLU(), nn.Linear(embed_dim // 2, embed_dim) ) elif mode == 'none': pass else: raise ValueError("mode must be 'sinusoidal' | 'learnable' | 'time_fourier' | 'time_learnable' | 'none'") def forward(self, x, times: torch.Tensor = None): if self.mode == 'none': return x B, L, D = x.shape if self.mode == 'learnable': return x + self.pos_embed[:, :L, :].to(x.dtype).to(x.device) if self.mode == 'sinusoidal': return x + self.pos_embed_buf[:, :L, :].to(x.dtype).to(x.device) if self.mode == 'time_fourier': if times is None: raise ValueError("times required for mode='time_fourier'") if times.dim() == 1: times = times.unsqueeze(0).expand(B, -1) x_t = times.unsqueeze(-1) * self.time_omega.view(1, 1, -1).to(times.device) pe = torch.cat([torch.sin(x_t), torch.cos(x_t)], dim=-1) pe = self.time_proj(pe.to(x.dtype)) return x + pe if self.mode == 'time_learnable': if times is None: raise ValueError("times required for mode='time_learnable'") if times.dim() == 1: times = times.unsqueeze(0).expand(B, -1) t0 = self.time_min_s t1 = self.time_max_s t_norm = (times - t0) / max(t1 - t0, self.time_eps) pe = self.time_mlp(t_norm.unsqueeze(-1).to(x.dtype)) return x + pe return x # ---------- Patch Embedding ---------- class PatchEmbedding(nn.Module): def __init__(self, in_channels, patch_size, patch_stride, embed_dim, temporal_dim, activation="gelu"): super().__init__() mid_dim = int(embed_dim / 2) dropout = 0.3 self.patch_size = patch_size self.patch_stride = patch_stride self.temporal_EEG = nn.Conv2d( in_channels=1, out_channels=temporal_dim, kernel_size=(1, patch_size), stride=(1, patch_stride)) self.spatial = nn.Conv2d( in_channels=temporal_dim, out_channels=temporal_dim, kernel_size=(in_channels, 1), stride=(1, 1)) self.temporal_PPG = nn.Conv1d( in_channels=1, out_channels=temporal_dim, kernel_size=patch_size, stride=patch_stride) self.proj = nn.Conv1d( in_channels=temporal_dim, out_channels=embed_dim, kernel_size=1, stride=1) self.dropout = nn.Dropout(dropout) self.norm = nn.LayerNorm(mid_dim) if activation == 'GELU': self.act = nn.GELU() elif activation == 'ReLU': self.act = nn.ReLU() elif activation == 'ELU': self.act = nn.ELU() def forward(self, x): # x: (B, 2, L) if x.shape[1] == 2: x = x.unsqueeze(1) # (B,1,2,L) x = self.temporal_EEG(x) x = self.spatial(x) x = x.squeeze(2) # (B,spatial_dim,L) else: x = self.temporal_PPG(x) # (B,embed_dim,Lp) x = self.act(x) x = self.dropout(x) x = self.proj(x) x = x.permute(0, 2, 1) return x # (B,Lp,embed_dim) # ---------- Self-Attention Block ---------- class SelfAttnBlock(nn.Module): def __init__(self, dim, heads=4, mlp_ratio=4, dropout=0.5, drop_prob=0.5, isPreLN=True): super().__init__() self.isPreLN = isPreLN self.attn = nn.MultiheadAttention(dim, heads, batch_first=True, dropout=dropout) self.ln1 = nn.LayerNorm(dim) self.ffn = nn.Sequential( nn.Linear(dim, dim * mlp_ratio), nn.ReLU(), nn.Dropout(dropout), nn.Linear(dim * mlp_ratio, dim), nn.Dropout(dropout), ) self.ln2 = nn.LayerNorm(dim) self.last_attn = None self.last_attn_out = None self.drop_path = DropPath(drop_prob) def forward(self, x): if self.isPreLN: x_norm = self.ln1(x) attn_out, attn = self.attn(x_norm, x_norm, x_norm, average_attn_weights=False) self.last_attn_out = attn_out if attn_out.requires_grad: self.last_attn_out.retain_grad() self.last_attn = attn if attn.requires_grad: self.last_attn.retain_grad() x = x + self.drop_path(attn_out) x_norm2 = self.ln2(x) f = self.ffn(x_norm2) x = x + self.drop_path(f) return x, attn_out, attn else: attn_out, attn = self.attn(x, x, x, average_attn_weights=False) self.last_attn_out = attn_out if attn_out.requires_grad: self.last_attn_out.retain_grad() self.last_attn = attn if attn.requires_grad: self.last_attn.retain_grad() x = self.ln1(x + attn_out) f = self.ffn(x) x = self.ln2(x + f) return x, attn_out, attn # ---------- Cross-Attention Block ---------- class CrossAttnBlock(nn.Module): def __init__(self, dim, heads=4, mlp_ratio=4, dropout=0.5, drop_prob=0.5, isPreLN=True): super().__init__() self.isPreLN = isPreLN self.attn = nn.MultiheadAttention(dim, heads, batch_first=True, dropout=dropout) self.ln1 = nn.LayerNorm(dim) self.ffn = nn.Sequential( nn.Linear(dim, dim * mlp_ratio), nn.ReLU(), nn.Dropout(dropout), nn.Linear(dim * mlp_ratio, dim), nn.Dropout(dropout), ) self.ln2 = nn.LayerNorm(dim) self.last_attn = None self.last_attn_out = None self.drop_path = DropPath(drop_prob) def forward(self, q, kv): if self.isPreLN: q_norm = self.ln1(q) attn_out, attn = self.attn(q_norm, kv, kv, average_attn_weights=False) self.last_attn_out = attn_out if attn_out.requires_grad: self.last_attn_out.retain_grad() self.last_attn = attn if attn.requires_grad: self.last_attn.retain_grad() x = q + self.drop_path(attn_out) x_norm = self.ln2(x) f = self.ffn(x_norm) x = x + self.drop_path(f) return x, attn_out, attn else: attn_out, attn = self.attn(q, kv, kv, average_attn_weights=False) self.last_attn_out = attn_out if attn_out.requires_grad: self.last_attn_out.retain_grad() self.last_attn = attn if attn.requires_grad: self.last_attn.retain_grad() x = self.ln1(q + attn_out) f = self.ffn(x) x = self.ln2(x + f) return x, attn_out, attn # ---------- BHITransformer (Parallel + Ablation) ---------- class BHITransformer(nn.Module): def __init__(self, eeg_channels, ppg_channels, eeg_patch_size, ppg_patch_size, eeg_patch_stride, ppg_patch_stride, embed_dim, temporal_dim, activation, num_heads, num_layers=1, num_classes=2, pos_encoding_mode='time_learnable', pe_max_len=4096, time_n_freq=16, time_max_freq=10, time_min_s=0.0, time_max_s=8.0, dropout=0.3, droppath=0.5, isPreLN=True, ablation_mode='Full'): super().__init__() self.valid_modes = ['Full', 'Self', 'Cross', 'Self_EEG', 'Self_PPG', 'Cross_E2P', 'Cross_P2E'] if ablation_mode not in self.valid_modes: raise ValueError(f"Unknown ablation_mode: {ablation_mode}. Choose from {self.valid_modes}") self.ablation_mode = ablation_mode # 1. Patch Embedding self.eeg_embed = PatchEmbedding(eeg_channels, eeg_patch_size, eeg_patch_stride, embed_dim, temporal_dim, activation) self.ppg_embed = PatchEmbedding(ppg_channels, ppg_patch_size, ppg_patch_stride, embed_dim, temporal_dim, activation) # 2. Positional Encoding self.pos_mode = pos_encoding_mode self.eeg_pos = PositionalEncoding(embed_dim, mode=pos_encoding_mode, max_len=pe_max_len, time_n_freq=time_n_freq, time_max_freq=time_max_freq, time_min_s=time_min_s, time_max_s=time_max_s) self.ppg_pos = PositionalEncoding(embed_dim, mode=pos_encoding_mode, max_len=pe_max_len, time_n_freq=time_n_freq, time_max_freq=time_max_freq, time_min_s=time_min_s, time_max_s=time_max_s) # 3. Transformer Branches self.self_attn_eeg = nn.ModuleList([SelfAttnBlock(dim=embed_dim, heads=num_heads, mlp_ratio=4, dropout=dropout, drop_prob=droppath, isPreLN=isPreLN) for _ in range(num_layers)]) self.self_attn_ppg = nn.ModuleList([SelfAttnBlock(dim=embed_dim, heads=num_heads, mlp_ratio=4, dropout=dropout, drop_prob=droppath, isPreLN=isPreLN) for _ in range(num_layers)]) self.cross_attn_e2p = nn.ModuleList([CrossAttnBlock(dim=embed_dim, heads=num_heads, mlp_ratio=4, dropout=dropout, drop_prob=droppath, isPreLN=isPreLN) for _ in range(num_layers)]) self.cross_attn_p2e = nn.ModuleList([CrossAttnBlock(dim=embed_dim, heads=num_heads, mlp_ratio=4, dropout=dropout, drop_prob=droppath, isPreLN=isPreLN) for _ in range(num_layers)]) # 4. Ablation 동적 차원 할당 if self.ablation_mode == 'Full': input_dim = 8 * embed_dim elif self.ablation_mode in ['Self', 'Cross']: input_dim = 4 * embed_dim else: input_dim = 2 * embed_dim bottleneck_dim = embed_dim self.classifier = nn.Sequential( nn.LayerNorm(input_dim), nn.Linear(input_dim, bottleneck_dim), nn.GELU(), nn.Dropout(dropout + 0.2), nn.Linear(bottleneck_dim, num_classes) ) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) if 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 _get_bhi_summary(self, x): avg_pool = x.mean(dim=1) max_pool, _ = x.max(dim=1) return torch.cat([avg_pool, max_pool], dim=-1) def forward(self, eeg, ppg, eeg_times=None, ppg_times=None): eeg_tok = self.eeg_embed(eeg) ppg_tok = self.ppg_embed(ppg) fs = 256.0 B, Le, _ = eeg_tok.shape B, Lp, _ = ppg_tok.shape eeg_dyn_time = (torch.arange(Le, device=eeg.device) * self.eeg_embed.patch_stride + (self.eeg_embed.patch_size / 2)) / fs ppg_dyn_time = (torch.arange(Lp, device=ppg.device) * self.ppg_embed.patch_stride + (self.ppg_embed.patch_size / 2)) / fs eeg_times = eeg_dyn_time.unsqueeze(0).expand(B, -1) ppg_times = ppg_dyn_time.unsqueeze(0).expand(B, -1) if self.pos_mode.startswith('time'): eeg_feat = self.eeg_pos(eeg_tok, times=eeg_times) ppg_feat = self.ppg_pos(ppg_tok, times=ppg_times) else: eeg_feat = self.eeg_pos(eeg_tok) ppg_feat = self.ppg_pos(ppg_tok) final_vectors = [] if self.ablation_mode in ['Full', 'Self', 'Self_EEG']: self_eeg = eeg_feat for blk in self.self_attn_eeg: self_eeg, _, _ = blk(self_eeg) final_vectors.append(self._get_bhi_summary(self_eeg)) if self.ablation_mode in ['Full', 'Self', 'Self_PPG']: self_ppg = ppg_feat for blk in self.self_attn_ppg: self_ppg, _, _ = blk(self_ppg) final_vectors.append(self._get_bhi_summary(self_ppg)) if self.ablation_mode in ['Full', 'Cross', 'Cross_E2P']: cross_e2p = eeg_feat for blk in self.cross_attn_e2p: cross_e2p, _, _ = blk(cross_e2p, ppg_feat) final_vectors.append(self._get_bhi_summary(cross_e2p)) if self.ablation_mode in ['Full', 'Cross', 'Cross_P2E']: cross_p2e = ppg_feat for blk in self.cross_attn_p2e: cross_p2e, _, _ = blk(cross_p2e, eeg_feat) final_vectors.append(self._get_bhi_summary(cross_p2e)) final_vector = torch.cat(final_vectors, dim=-1) logits = self.classifier(final_vector) return logits