import os import torch import torch.nn as nn import torch.nn.functional as F import pytorch_lightning as pl from torch.utils.data import Dataset, DataLoader from typing import List, Tuple, Optional import numpy as np from pathlib import Path from deepspeed.ops.adam import FusedAdam class MusicClassifier(pl.LightningModule): def __init__(self, input_dim: int, hidden_dim: int = 256, learning_rate: float = 1e-4, emb_model: Optional[nn.Module] = None, is_emb: bool = False): super().__init__() self.save_hyperparameters() self.model = SegmentTransformer( input_dim=input_dim, hidden_dim=hidden_dim ) self.emb_model = emb_model self.learning_rate = learning_rate self.is_emb = is_emb def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: return self.model(x, mask) def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor: x, y, mask = batch if self.is_emb == False: _,x = self.emb_model(x) y_hat = self(x, mask) loss = F.binary_cross_entropy_with_logits(y_hat.squeeze(), y.float()) # Log metrics self.log('train_loss', loss,on_epoch=True, prog_bar=True) preds = (torch.sigmoid(y_hat.squeeze()) > 0.5).float() acc = (preds == y.float()).float().mean() self.log('train_acc', acc,on_epoch=True, prog_bar=True) return loss def validation_step(self, batch: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int) -> None: x, y, mask = batch if self.is_emb == False: _,x = self.emb_model(x) y_hat = self(x, mask) loss = F.binary_cross_entropy_with_logits(y_hat.squeeze(), y.float()) # Calculate accuracy preds = (torch.sigmoid(y_hat.squeeze()) > 0.5).float() acc = (preds == y.float()).float().mean() self.log('val_loss', loss, prog_bar=True) self.log('val_acc', acc, prog_bar=True) def configure_optimizers(self): optimizer = torch.optim.AdamW( self.parameters(), lr=self.learning_rate, weight_decay=0.01 ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=100, # Adjust based on your training epochs eta_min=1e-6 ) return { 'optimizer': optimizer, 'lr_scheduler': scheduler, 'monitor': 'val_loss' } class MusicAudioClassifier(pl.LightningModule): def __init__(self, input_dim: int, hidden_dim: int = 256, learning_rate: float = 1e-4, emb_model: Optional[nn.Module] = None, is_emb: bool = False, mode: str = 'only_emb', share_parameter: bool = False): super().__init__() self.save_hyperparameters() self.model = SegmentTransformer( input_dim=input_dim, hidden_dim=hidden_dim, mode = mode, share_parameter = share_parameter ) self.emb_model = emb_model self.learning_rate = learning_rate self.is_emb = is_emb def _process_audio_batch(self, x: torch.Tensor) -> torch.Tensor: B, S = x.shape[:2] # [B, S, C, M, T] or [B, S, C, T] for wav, [B, S, 1?, embsize] for emb x = x.view(B*S, *x.shape[2:]) # [B*S, C, M, T] if self.is_emb == False: _, embeddings = self.emb_model(x) # [B*S, emb_dim] else: embeddings = x if embeddings.dim() == 3: pooled_features = embeddings.mean(dim=1) # transformer else: pooled_features = embeddings # CCV..? no need to pooling return pooled_features.view(B, S, -1) # [B, S, emb_dim] def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: x = self._process_audio_batch(x) # 이걸 freeze하고 쓰는게 사실상 윗버전임 x = x.half() return self.model(x, mask) def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor: x, y, mask = batch x = x.half() y_hat = self(x, mask) # 배치 크기가 1인 경우 예외처리 if y_hat.size(0) == 1: loss = F.binary_cross_entropy_with_logits(y_hat.flatten(), y.float().flatten()) probs = torch.sigmoid(y_hat.flatten()) y_true = y.float().flatten() else: loss = F.binary_cross_entropy_with_logits(y_hat.squeeze(), y.float()) probs = torch.sigmoid(y_hat.squeeze()) y_true = y.float() # 간단한 배치 손실만 로깅 (step 수준) self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True) # 전체 에폭에 대한 메트릭 계산을 위해 예측과 실제값 저장 self.training_step_outputs.append({'preds': probs, 'targets': y_true}) return loss def validation_step(self, batch: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int) -> None: x, y, mask = batch x = x.half() y_hat = self(x, mask) # 배치 크기가 1인 경우 예외처리 if y_hat.size(0) == 1: loss = F.binary_cross_entropy_with_logits(y_hat.flatten(), y.float().flatten()) probs = torch.sigmoid(y_hat.flatten()) y_true = y.float().flatten() else: loss = F.binary_cross_entropy_with_logits(y_hat.squeeze(), y.float()) probs = torch.sigmoid(y_hat.squeeze()) y_true = y.float() # 간단한 배치 손실만 로깅 (step 수준) self.log('val_loss', loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True) # 전체 에폭에 대한 메트릭 계산을 위해 예측과 실제값 저장 self.validation_step_outputs.append({'preds': probs, 'targets': y_true}) def test_step(self, batch: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int) -> None: x, y, mask = batch x = x.half() y_hat = self(x, mask) # 배치 크기가 1인 경우 예외처리 if y_hat.size(0) == 1: loss = F.binary_cross_entropy_with_logits(y_hat.flatten(), y.float().flatten()) probs = torch.sigmoid(y_hat.flatten()) y_true = y.float().flatten() else: loss = F.binary_cross_entropy_with_logits(y_hat.squeeze(), y.float()) probs = torch.sigmoid(y_hat.squeeze()) y_true = y.float() # 간단한 배치 손실만 로깅 (step 수준) self.log('test_loss', loss, on_epoch=True, prog_bar=True) # 전체 에폭에 대한 메트릭 계산을 위해 예측과 실제값 저장 self.test_step_outputs.append({'preds': probs, 'targets': y_true}) def on_train_epoch_start(self): # 에폭 시작 시 결과 저장용 리스트 초기화 self.training_step_outputs = [] def on_validation_epoch_start(self): # 에폭 시작 시 결과 저장용 리스트 초기화 self.validation_step_outputs = [] def on_test_epoch_start(self): # 에폭 시작 시 결과 저장용 리스트 초기화 self.test_step_outputs = [] def on_train_epoch_end(self): # 에폭이 끝날 때 전체 데이터에 대한 메트릭 계산 if not hasattr(self, 'training_step_outputs') or not self.training_step_outputs: return all_preds = torch.cat([x['preds'] for x in self.training_step_outputs]) all_targets = torch.cat([x['targets'] for x in self.training_step_outputs]) # 전체 데이터에 대한 메트릭 계산 binary_preds = (all_preds > 0.5).float() # 정확도 계산 acc = (binary_preds == all_targets).float().mean() # 혼동 행렬 요소 계산 tp = torch.sum((binary_preds == 1) & (all_targets == 1)).float() fp = torch.sum((binary_preds == 1) & (all_targets == 0)).float() tn = torch.sum((binary_preds == 0) & (all_targets == 0)).float() fn = torch.sum((binary_preds == 0) & (all_targets == 1)).float() # 메트릭 계산 precision = tp / (tp + fp) if (tp + fp) > 0 else torch.tensor(0.0).to(tp.device) recall = tp / (tp + fn) if (tp + fn) > 0 else torch.tensor(0.0).to(tp.device) f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else torch.tensor(0.0).to(tp.device) specificity = tn / (tn + fp) if (tn + fp) > 0 else torch.tensor(0.0).to(tn.device) # 로깅 - 일관된 이름 사용 self.log('train_acc', acc, on_epoch=True, prog_bar=True, sync_dist=True) self.log('train_precision', precision, on_epoch=True, sync_dist=True) self.log('train_recall', recall, on_epoch=True, sync_dist=True) self.log('train_f1', f1, on_epoch=True, prog_bar=True, sync_dist=True) self.log('train_specificity', specificity, on_epoch=True, sync_dist=True) def on_validation_epoch_end(self): # 에폭이 끝날 때 전체 데이터에 대한 메트릭 계산 if not hasattr(self, 'validation_step_outputs') or not self.validation_step_outputs: return all_preds = torch.cat([x['preds'] for x in self.validation_step_outputs]) all_targets = torch.cat([x['targets'] for x in self.validation_step_outputs]) # ROC-AUC 계산 (간단한 근사) sorted_indices = torch.argsort(all_preds, descending=True) sorted_targets = all_targets[sorted_indices] n_pos = torch.sum(all_targets) n_neg = len(all_targets) - n_pos if n_pos > 0 and n_neg > 0: # TPR과 FPR을 누적합으로 계산 tpr_curve = torch.cumsum(sorted_targets, dim=0) / n_pos fpr_curve = torch.cumsum(1 - sorted_targets, dim=0) / n_neg # AUC 계산 (사다리꼴 법칙) width = fpr_curve[1:] - fpr_curve[:-1] height = (tpr_curve[1:] + tpr_curve[:-1]) / 2 auc_approx = torch.sum(width * height) self.log('val_auc', auc_approx, on_epoch=True) # 전체 데이터에 대한 메트릭 계산 binary_preds = (all_preds > 0.5).float() # 정확도 계산 acc = (binary_preds == all_targets).float().mean() # 혼동 행렬 요소 계산 tp = torch.sum((binary_preds == 1) & (all_targets == 1)).float() fp = torch.sum((binary_preds == 1) & (all_targets == 0)).float() tn = torch.sum((binary_preds == 0) & (all_targets == 0)).float() fn = torch.sum((binary_preds == 0) & (all_targets == 1)).float() # 메트릭 계산 precision = tp / (tp + fp) if (tp + fp) > 0 else torch.tensor(0.0).to(tp.device) recall = tp / (tp + fn) if (tp + fn) > 0 else torch.tensor(0.0).to(tp.device) f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else torch.tensor(0.0).to(tp.device) specificity = tn / (tn + fp) if (tn + fp) > 0 else torch.tensor(0.0).to(tn.device) # 로깅 - 일관된 이름 사용 (val_epoch_f1 대신 val_f1 사용) self.log('val_acc', acc, on_epoch=True, prog_bar=True, sync_dist=True) self.log('val_precision', precision, on_epoch=True, sync_dist=True) self.log('val_recall', recall, on_epoch=True, sync_dist=True) self.log('val_f1', f1, on_epoch=True, prog_bar=True, sync_dist=True) self.log('val_specificity', specificity, on_epoch=True, sync_dist=True) def on_test_epoch_end(self): # 에폭이 끝날 때 전체 테스트 데이터에 대한 메트릭 계산 if not hasattr(self, 'test_step_outputs') or not self.test_step_outputs: return all_preds = torch.cat([x['preds'] for x in self.test_step_outputs]) all_targets = torch.cat([x['targets'] for x in self.test_step_outputs]) # ROC-AUC 계산 (간단한 근사) sorted_indices = torch.argsort(all_preds, descending=True) sorted_targets = all_targets[sorted_indices] n_pos = torch.sum(all_targets) n_neg = len(all_targets) - n_pos if n_pos > 0 and n_neg > 0: # TPR과 FPR을 누적합으로 계산 tpr_curve = torch.cumsum(sorted_targets, dim=0) / n_pos fpr_curve = torch.cumsum(1 - sorted_targets, dim=0) / n_neg # AUC 계산 (사다리꼴 법칙) width = fpr_curve[1:] - fpr_curve[:-1] height = (tpr_curve[1:] + tpr_curve[:-1]) / 2 auc_approx = torch.sum(width * height) self.log('test_auc', auc_approx, on_epoch=True, sync_dist=True) # 전체 데이터에 대한 메트릭 계산 binary_preds = (all_preds > 0.5).float() # 정확도 계산 acc = (binary_preds == all_targets).float().mean() # 혼동 행렬 요소 계산 tp = torch.sum((binary_preds == 1) & (all_targets == 1)).float() fp = torch.sum((binary_preds == 1) & (all_targets == 0)).float() tn = torch.sum((binary_preds == 0) & (all_targets == 0)).float() fn = torch.sum((binary_preds == 0) & (all_targets == 1)).float() # 메트릭 계산 precision = tp / (tp + fp) if (tp + fp) > 0 else torch.tensor(0.0).to(tp.device) recall = tp / (tp + fn) if (tp + fn) > 0 else torch.tensor(0.0).to(tp.device) f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else torch.tensor(0.0).to(tp.device) specificity = tn / (tn + fp) if (tn + fp) > 0 else torch.tensor(0.0).to(tn.device) balanced_acc = (recall + specificity) / 2 # 로깅 - 일관된 이름 사용 self.log('test_acc', acc, on_epoch=True, prog_bar=True) self.log('test_precision', precision, on_epoch=True) self.log('test_recall', recall, on_epoch=True) self.log('test_f1', f1, on_epoch=True, prog_bar=True) self.log('test_specificity', specificity, on_epoch=True) self.log('test_balanced_acc', balanced_acc, on_epoch=True) def configure_optimizers(self): optimizer = FusedAdam(self.parameters(),lr=self.learning_rate, weight_decay=0.01) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=100, # Adjust based on your training epochs eta_min=1e-6 ) return { 'optimizer': optimizer, 'lr_scheduler': scheduler, 'monitor': 'val_loss', } def pad_sequence_with_mask(batch: List[Tuple[torch.Tensor, int]]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Collate function for DataLoader that creates padded sequences and attention masks with fixed length (48).""" embeddings, labels = zip(*batch) fixed_len = 48 # 고정 길이 batch_size = len(embeddings) feat_dim = embeddings[0].shape[-1] padded = torch.zeros((batch_size, fixed_len, feat_dim)) # 고정 길이로 패딩된 텐서 mask = torch.ones((batch_size, fixed_len), dtype=torch.bool) # True는 padding을 의미 for i, emb in enumerate(embeddings): length = emb.shape[0] # 길이가 고정 길이보다 길면 자르고, 짧으면 패딩 if length > fixed_len: padded[i, :] = emb[:fixed_len] # fixed_len보다 긴 부분을 잘라서 채운다. mask[i, :] = False else: padded[i, :length] = emb # 실제 데이터 길이에 맞게 채운다. mask[i, :length] = False # 패딩이 아닌 부분은 False로 설정 return padded, torch.tensor(labels), mask class SegmentTransformer(nn.Module): def __init__(self, input_dim: int, hidden_dim: int = 256, num_heads: int = 8, num_layers: int = 4, dropout: float = 0.1, max_sequence_length: int = 1000, mode: str = 'only_emb', share_parameter: bool = False): super().__init__() # Original sequence processing self.input_projection = nn.Linear(input_dim, hidden_dim) self.mode = mode self.share_parameter = share_parameter # Positional encoding position = torch.arange(max_sequence_length).unsqueeze(1) div_term = torch.exp(torch.arange(0, hidden_dim, 2) * (-np.log(10000.0) / hidden_dim)) pos_encoding = torch.zeros(max_sequence_length, hidden_dim) pos_encoding[:, 0::2] = torch.sin(position * div_term) pos_encoding[:, 1::2] = torch.cos(position * div_term) self.register_buffer('pos_encoding', pos_encoding) # Transformer for original sequence encoder_layer = nn.TransformerEncoderLayer( d_model=hidden_dim, nhead=num_heads, dim_feedforward=hidden_dim * 4, dropout=dropout, batch_first=True ) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) self.sim_transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) # Self-similarity stream processing self.similarity_projection = nn.Sequential( nn.Conv1d(1, hidden_dim // 2, kernel_size=3, padding=1), nn.ReLU(), nn.Conv1d(hidden_dim // 2, hidden_dim, kernel_size=3, padding=1), nn.ReLU(), nn.Dropout(dropout) ) # Transformer for similarity stream self.similarity_transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) # Final classification head self.classification_head_dim = hidden_dim * 2 if mode == 'both' else hidden_dim self.classification_head = nn.Sequential( nn.Linear(self.classification_head_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_dim, hidden_dim // 2), nn.LayerNorm(hidden_dim // 2), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_dim // 2, 1) ) def forward(self, x: torch.Tensor, padding_mask: Optional[torch.Tensor] = None) -> torch.Tensor: batch_size, seq_len, _ = x.shape # 1. Process original sequence x1 = self.input_projection(x) x1 = x1 + self.pos_encoding[:seq_len].unsqueeze(0) x1 = self.transformer(x1, src_key_padding_mask=padding_mask) # padding_mask 사용 # 2. Calculate and process self-similarity x_expanded = x.unsqueeze(2) x_transposed = x.unsqueeze(1) distances = torch.mean((x_expanded - x_transposed) ** 2, dim=-1) similarity_matrix = torch.exp(-distances) # (batch_size, seq_len, seq_len) # 자기 유사도 마스크 생성 및 적용 (각 시점에 대한 마스크 개별 적용) if padding_mask is not None: similarity_mask = padding_mask.unsqueeze(1) | padding_mask.unsqueeze(2) # (batch_size, seq_len, seq_len) similarity_matrix = similarity_matrix.masked_fill(similarity_mask, 0.0) # Process similarity matrix row by row using Conv1d x2 = similarity_matrix.unsqueeze(1) # (batch_size, 1, seq_len, seq_len) x2 = x2.view(batch_size * seq_len, 1, seq_len) # Reshape for Conv1d x2 = self.similarity_projection(x2) # (batch_size * seq_len, hidden_dim, seq_len) x2 = x2.mean(dim=2) # Pool across sequence dimension x2 = x2.view(batch_size, seq_len, -1) # Reshape back x2 = x2 + self.pos_encoding[:seq_len].unsqueeze(0) if self.share_parameter: x2 = self.transformer(x2, src_key_padding_mask=padding_mask) else: x2 = self.sim_transformer(x2, src_key_padding_mask=padding_mask) # padding_mask 사용 # 3. Global average pooling for both streams if padding_mask is not None: mask_expanded = (~padding_mask).float().unsqueeze(-1) x1 = (x1 * mask_expanded).sum(dim=1) / mask_expanded.sum(dim=1) x2 = (x2 * mask_expanded).sum(dim=1) / mask_expanded.sum(dim=1) else: x1 = x1.mean(dim=1) x2 = x2.mean(dim=1) # 4. Combine both streams and classify #x = x1 # only emb #x = x2 # only structure #x = torch.cat([x1, x2], dim=-1) if self.mode == 'only_emb': x = x1 elif self.mode == 'only_structure': x = x2 elif self.mode == 'both': x = torch.cat([x1, x2], dim=-1) x = x.half() return self.classification_head(x)