Armak/SED
0
1import os2import random3import warnings4from pathlib import Path5from copy import deepcopy6 7import numpy as np8import pandas as pd9import pytorch_lightning as pl10import sed_scores_eval11import torch12import torchmetrics13import yaml14from torchaudio.transforms import AmplitudeToDB, MelSpectrogram15 16from desed_task.data_augm import mixup, filt_aug_prototype, freq_mask17from desed_task.evaluation.evaluation_measures import (18 compute_per_intersection_macro_f1,19 compute_psds_from_operating_points,20 compute_psds_from_scores,21)22from desed_task.utils.scaler import TorchScaler23from encoder import ManyHotEncoder24from model import CRNN25from utils import (26 batched_decode_preds,27 classes_labels,28 log_sedeval_metrics,29)30 31 32class SED(pl.LightningModule):33 34 """Pytorch lightning module for the SED task.35 Args:36 hparams: dict, the dictionary to be used for the current experiment.37 encoder: ManyHotEncoder object, object to encode and decode labels.38 sed: torch.Module, the model to be trained.39 opt: torch.optimizer.Optimizer object, the optimizer to be used.40 train_data: torch.utils.data.Dataset subclass object, the training data to be used.41 valid_data: torch.utils.data.Dataset subclass object, the validation data to be used.42 test_data: torch.utils.data.Dataset subclass object, the test data to be used.43 train_sampler: torch.utils.data.Sampler subclass object, the sampler to be used in the training dataloader.44 scheduler: BaseScheduler subclass object, the scheduler to be used.45 This is used to apply ramp-up during training for example.46 fast_dev_run: bool, whether to launch a run with only one batch for each set, this is for development purpose,47 to test the code runs.48 """49 50 def __init__(51 self,52 hparams,53 encoder,54 sed,55 opt=None,56 train_data=None,57 valid_data=None,58 test_data=None,59 train_sampler=None,60 scheduler=None,61 fast_dev_run=False,62 evaluation=False,63 ):64 super(SED, self).__init__()65 self.hparams.update(hparams)66 67 self.encoder = encoder68 self.sed_student = sed69 self.sed_teacher = deepcopy(sed)70 self.opt = opt71 self.train_data = train_data72 self.valid_data = valid_data73 self.test_data = test_data74 self.train_sampler = train_sampler75 self.scheduler = scheduler76 self.fast_dev_run = fast_dev_run77 self.evaluation = evaluation78 79 if self.fast_dev_run:80 self.num_workers = 181 else:82 self.num_workers = self.hparams["training"]["num_workers"]83 84 feat_params = self.hparams["feats"]85 86 self.mel_spec = MelSpectrogram(87 sample_rate=feat_params["sample_rate"],88 n_fft=feat_params["n_window"],89 win_length=feat_params["n_window"],90 hop_length=feat_params["hop_length"],91 f_min=feat_params["f_min"],92 f_max=feat_params["f_max"],93 n_mels=feat_params["n_mels"],94 window_fn=torch.hamming_window,95 wkwargs={"periodic": False},96 power=1,97 center=False,98 )99 100 for param in self.sed_teacher.parameters():101 param.detach_()102 103 # * instantiating loss fns and scaler104 self.supervised_loss = torch.nn.BCELoss()105 106 if hparams["training"]["consistency_loss"] == "mse":107 self.consistency_loss = torch.nn.MSELoss()108 elif hparams["training"]["consistency_loss"] == "bce":109 self.consistency_loss = torch.nn.BCELoss()110 else:111 raise NotImplementedError112 113 self.get_weak_student_f1_seg_macro = (114 torchmetrics.classification.f_beta.MultilabelF1Score(115 len(self.encoder.labels),116 average="macro",117 )118 )119 120 self.get_weak_teacher_f1_seg_macro = (121 torchmetrics.classification.f_beta.MultilabelF1Score(122 len(self.encoder.labels),123 average="macro",124 )125 )126 127 self.scaler = self._init_scaler()128 129 # * buffer for event based scores which we compute using sed-eval130 self.val_buffer_student_strong = {131 k: pd.DataFrame() for k in self.hparams["training"]["val_thresholds"]132 }133 134 self.val_buffer_student_test = {135 k: pd.DataFrame() for k in self.hparams["training"]["val_thresholds"]136 }137 138 self.val_buffer_teacher_strong = {139 k: pd.DataFrame() for k in self.hparams["training"]["val_thresholds"]140 }141 142 self.val_buffer_teacher_test = {143 k: pd.DataFrame() for k in self.hparams["training"]["val_thresholds"]144 }145 146 self.val_scores_postprocessed_buffer_student_strong = {}147 self.val_scores_postprocessed_buffer_teacher_strong = {}148 149 test_n_thresholds = self.hparams["training"]["n_test_thresholds"]150 test_thresholds = np.arange(151 1 / (test_n_thresholds * 2), 1, 1 / test_n_thresholds152 )153 self.test_psds_buffer_student = {k: pd.DataFrame() for k in test_thresholds}154 self.test_psds_buffer_teacher = {k: pd.DataFrame() for k in test_thresholds}155 self.decoded_student_05_buffer = pd.DataFrame()156 self.decoded_teacher_05_buffer = pd.DataFrame()157 self.test_scores_raw_buffer_student = {}158 self.test_scores_raw_buffer_teacher = {}159 self.test_scores_postprocessed_buffer_student = {}160 self.test_scores_postprocessed_buffer_teacher = {}161 162 _exp_dir = None163 164 @property165 def exp_dir(self):166 if self._exp_dir is None:167 try:168 self._exp_dir = self.logger.log_dir169 if self._exp_dir is None:170 self._exp_dir = self.hparams["log_dir"]171 except Exception:172 self._exp_dir = self.hparams["log_dir"]173 return self._exp_dir174 175 def lr_scheduler_step(self, scheduler, optimizer_idx, metric):176 scheduler.step()177 178 def on_train_start(self) -> None:179 if not self.fast_dev_run:180 to_ignore = [181 ".*Trying to infer the `batch_size` from an ambiguous collection.*",182 ".*invalid value encountered in divide*",183 ".*mean of empty slice*",184 ".*self.log*",185 ]186 for message in to_ignore:187 warnings.filterwarnings("ignore", message)188 189 def update_ema(self, alpha, global_step, student_model, teacher_model):190 """Update teacher model parameters191 192 Args:193 alpha: float, the factor to be used between each updated step.194 global_step: int, the current global step to be used.195 student_model: torch.Module, student model to use196 teacher_model: torch.Module, teacher model to use197 """198 # Use the true average until the exponential average is more correct199 alpha = min(1 - 1 / (global_step + 1), alpha)200 for teacher_params, student_params in zip(201 teacher_model.parameters(), student_model.parameters()202 ):203 teacher_params.data.mul_(alpha).add_(student_params.data, alpha=1 - alpha)204 205 def _init_scaler(self):206 """Scaler inizialization function. It can be either a dataset or instance scaler.207 208 Raises:209 NotImplementedError: in case of not Implemented scaler210 211 Returns:212 TorchScaler: returns the scaler213 """214 215 if self.hparams["scaler"]["statistic"] == "instance":216 scaler = TorchScaler(217 "instance",218 self.hparams["scaler"]["normtype"],219 self.hparams["scaler"]["dims"],220 )221 return scaler222 elif self.hparams["scaler"]["statistic"] == "dataset":223 scaler = TorchScaler(224 "dataset",225 self.hparams["scaler"]["normtype"],226 self.hparams["scaler"]["dims"],227 )228 else:229 raise NotImplementedError230 231 if self.hparams["scaler"]["savepath"] is not None:232 if os.path.exists(self.hparams["scaler"]["savepath"]):233 scaler = torch.load(self.hparams["scaler"]["savepath"])234 print(235 "Loaded Scaler from previous checkpoint from {}".format(236 self.hparams["scaler"]["savepath"]237 )238 )239 return scaler240 241 self.train_loader = self.train_dataloader()242 scaler.fit(243 self.train_loader,244 transform_func=lambda x: self.take_log(self.mel_spec(x[0])),245 )246 247 if self.hparams["scaler"]["savepath"] is not None:248 torch.save(scaler, self.hparams["scaler"]["savepath"])249 print(250 "Saving Scaler from previous checkpoint at {}".format(251 self.hparams["scaler"]["savepath"]252 )253 )254 return scaler255 256 def take_log(self, mels):257 """Apply the log transformation to mel spectrograms.258 Args:259 mels: torch.Tensor, mel spectrograms for which to apply log.260 261 Returns:262 Tensor: logarithmic mel spectrogram of the mel spectrogram given as input263 """264 265 amp_to_db = AmplitudeToDB(stype="amplitude")266 amp_to_db.amin = 1e-5 # amin= 1e-5 as in librosa267 return amp_to_db(mels).clamp(min=-50, max=80)268 269 def training_step(self, batch, batch_indx):270 """Apply the training for one batch (a step). Used during trainer.fit271 272 Args:273 batch: torch.Tensor, batch input tensor274 batch_indx: torch.Tensor, 1D tensor of indexes to know which data are present in each batch.275 276 Returns:277 torch.Tensor, the loss to take into account.278 """279 280 indx_strong, indx_weak, _ = self.hparams["training"]["batch_size"]281 282 audio, labels, _ = batch283 features = self.mel_spec(audio)284 285 batch_num = features.shape[0]286 # deriving masks for each dataset287 strong_mask = torch.zeros(batch_num).to(features).bool()288 weak_mask = torch.zeros(batch_num).to(features).bool()289 strong_mask[:indx_strong] = 1290 weak_mask[indx_strong : indx_weak + indx_strong] = 1291 292 labels_weak = (torch.sum(labels[weak_mask], -1) > 0).float()293 294 mixup_type = self.hparams["training"].get("mixup")295 if mixup_type is not None and 0.5 > random.random():296 features[weak_mask], labels_weak = mixup(297 features[weak_mask], labels_weak, mixup_label_type=mixup_type298 )299 features[strong_mask], labels[strong_mask] = mixup(300 features[strong_mask], labels[strong_mask], mixup_label_type=mixup_type301 )302 303 features = self.take_log(features)304 305 if 0.1 > random.random():306 features[weak_mask] = filt_aug_prototype(features[weak_mask])307 features[strong_mask] = filt_aug_prototype(features[strong_mask])308 309 if 0.1 > random.random():310 features[weak_mask] = freq_mask(features[weak_mask])311 features[strong_mask] = freq_mask(features[strong_mask])312 313 strong_preds_student, weak_preds_student = self.sed_student(314 self.scaler(features)315 )316 317 loss_strong = self.supervised_loss(318 strong_preds_student[strong_mask], labels[strong_mask]319 )320 # supervised loss on weakly labelled321 loss_weak = self.supervised_loss(weak_preds_student[weak_mask], labels_weak)322 323 # total supervised loss324 total_supervised_loss = loss_strong + loss_weak325 326 with torch.no_grad():327 strong_preds_teacher, weak_preds_teacher = self.sed_teacher(328 self.scaler(features)329 )330 loss_strong_teacher = self.supervised_loss(331 strong_preds_teacher[strong_mask], labels[strong_mask]332 )333 334 loss_weak_teacher = self.supervised_loss(335 weak_preds_teacher[weak_mask], labels_weak336 )337 # we apply consistency between the predictions338 weight = (339 self.hparams["training"]["const_max"]340 * self.scheduler["scheduler"]._get_scaling_factor()341 )342 343 strong_consistency_loss = self.consistency_loss(344 strong_preds_student, strong_preds_teacher.detach()345 )346 weak_consistency_loss = self.consistency_loss(347 weak_preds_student, weak_preds_teacher.detach()348 )349 total_consistency_loss = (350 strong_consistency_loss + weak_consistency_loss351 ) * weight352 353 total_loss = total_supervised_loss + total_consistency_loss354 355 self.log(356 "train/student/loss_strong", loss_strong, prog_bar=True, sync_dist=True357 )358 self.log("train/student/loss_weak", loss_weak, prog_bar=True, sync_dist=True)359 self.log(360 "train/student/loss_total_supervised",361 total_supervised_loss,362 prog_bar=True,363 sync_dist=True,364 )365 self.log(366 "train/student/loss_consistency",367 total_consistency_loss,368 prog_bar=True,369 sync_dist=True,370 )371 self.log(372 "train/teacher/loss_strong",373 loss_strong_teacher,374 prog_bar=True,375 sync_dist=True,376 )377 self.log(378 "train/teacher/loss_weak", loss_weak_teacher, prog_bar=True, sync_dist=True379 )380 self.log("train/total_loss", total_loss, prog_bar=True, sync_dist=True)381 self.log(382 "train/step",383 self.scheduler["scheduler"].step_num,384 prog_bar=True,385 sync_dist=True,386 )387 self.log("train/lr", self.opt.param_groups[-1]["lr"], sync_dist=True)388 389 return total_loss390 391 def on_before_zero_grad(self, *args, **kwargs):392 # update EMA teacher393 self.update_ema(394 self.hparams["training"]["ema_factor"],395 self.scheduler["scheduler"].step_num,396 self.sed_student,397 self.sed_teacher,398 )399 400 def validation_step(self, batch, batch_indx):401 """Apply validation to a batch (step). Used during trainer.fit402 403 Args:404 batch: torch.Tensor, input batch tensor405 batch_indx: torch.Tensor, 1D tensor of indexes to know which data are present in each batch.406 Returns:407 """408 409 audio, labels, _, filenames = batch410 411 features = self.mel_spec(audio)412 strong_preds_student, weak_preds_student = self.sed_student(413 self.scaler(self.take_log(features))414 )415 strong_preds_teacher, weak_preds_teacher = self.sed_teacher(416 self.scaler(self.take_log(features))417 )418 419 weak_mask = (420 torch.tensor(421 [422 str(Path(x).parent)423 == str(Path(self.hparams["data"]["weak_folder"]))424 for x in filenames425 ]426 )427 .to(audio)428 .bool()429 )430 strong_mask = (431 torch.tensor(432 [433 str(Path(x).parent)434 == str(Path(self.hparams["data"]["synth_val_folder"]))435 for x in filenames436 ]437 )438 .to(audio)439 .bool()440 )441 442 if torch.any(weak_mask):443 labels_weak = (torch.sum(labels[weak_mask], -1) >= 1).float()444 loss_weak_student = self.supervised_loss(445 weak_preds_student[weak_mask], labels_weak446 )447 loss_weak_teacher = self.supervised_loss(448 weak_preds_teacher[weak_mask], labels_weak449 )450 self.log("val/weak/student/loss_weak", loss_weak_student)451 self.log("val/weak/teacher/loss_weak", loss_weak_teacher)452 453 # accumulate f1 score for weak labels454 self.get_weak_student_f1_seg_macro(455 weak_preds_student[weak_mask], labels_weak.long()456 )457 self.get_weak_teacher_f1_seg_macro(458 weak_preds_teacher[weak_mask], labels_weak.long()459 )460 461 if torch.any(strong_mask):462 loss_strong_student = self.supervised_loss(463 strong_preds_student[strong_mask], labels[strong_mask]464 )465 loss_strong_teacher = self.supervised_loss(466 strong_preds_teacher[strong_mask], labels[strong_mask]467 )468 469 self.log("val/strong/student/loss_strong", loss_strong_student)470 self.log("val/strong/teacher/loss_strong", loss_strong_teacher)471 472 filenames_strong = [473 x474 for x in filenames475 if Path(x).parent == Path(self.hparams["data"]["synth_val_folder"])476 ]477 478 (479 scores_raw_student_strong,480 scores_postprocessed_student_strong,481 decoded_student_strong,482 ) = batched_decode_preds(483 strong_preds_student[strong_mask],484 filenames_strong,485 self.encoder,486 median_filter=self.hparams["training"]["median_window"],487 thresholds=list(self.val_buffer_student_strong.keys()),488 )489 490 self.val_scores_postprocessed_buffer_student_strong.update(491 scores_postprocessed_student_strong492 )493 for th in self.val_buffer_student_strong.keys():494 self.val_buffer_student_strong[th] = pd.concat(495 [self.val_buffer_student_strong[th], decoded_student_strong[th]],496 ignore_index=True,497 )498 499 (500 scores_raw_teacher_strong,501 scores_postprocessed_teacher_strong,502 decoded_teacher_strong,503 ) = batched_decode_preds(504 strong_preds_teacher[strong_mask],505 filenames_strong,506 self.encoder,507 median_filter=self.hparams["training"]["median_window"],508 thresholds=list(self.val_buffer_teacher_strong.keys()),509 )510 511 self.val_scores_postprocessed_buffer_teacher_strong.update(512 scores_postprocessed_teacher_strong513 )514 for th in self.val_buffer_teacher_strong.keys():515 self.val_buffer_teacher_strong[th] = pd.concat(516 [self.val_buffer_teacher_strong[th], decoded_teacher_strong[th]],517 ignore_index=True,518 )519 520 # total supervised loss521 if torch.any(strong_mask) and torch.any(weak_mask):522 total_loss_student = loss_strong_student + loss_weak_student523 total_loss_teacher = loss_strong_teacher + loss_weak_teacher524 self.log(525 "val/student/total_loss",526 total_loss_student,527 prog_bar=True,528 sync_dist=True,529 )530 self.log(531 "val/teacher/total_loss",532 total_loss_teacher,533 prog_bar=True,534 sync_dist=True,535 )536 537 return538 539 def validation_epoch_end(self, outputs):540 """Function applied at the end of all the validation steps of the epoch.541 542 Args:543 outputs: torch.Tensor, the concatenation of everything returned by validation_step.544 545 Returns:546 torch.Tensor, the objective metric to be used to choose the best model from for example.547 """548 549 weak_student_f1_macro = self.get_weak_student_f1_seg_macro.compute()550 weak_teacher_f1_macro = self.get_weak_teacher_f1_seg_macro.compute()551 552 # * strong val dataset553 ground_truth = sed_scores_eval.io.read_ground_truth_events(554 self.hparams["data"]["synth_val_tsv"]555 )556 audio_durations = sed_scores_eval.io.read_audio_durations(557 self.hparams["data"]["synth_val_dur"]558 )559 if self.fast_dev_run:560 ground_truth = {561 audio_id: ground_truth[audio_id]562 for audio_id in self.val_scores_postprocessed_buffer_student_strong563 }564 audio_durations = {565 audio_id: audio_durations[audio_id]566 for audio_id in self.val_scores_postprocessed_buffer_student_strong567 }568 else:569 # * drop audios without events570 ground_truth = {571 audio_id: gt for audio_id, gt in ground_truth.items() if len(gt) > 0572 }573 audio_durations = {574 audio_id: audio_durations[audio_id] for audio_id in ground_truth.keys()575 }576 psds1_sed_scores_eval_student = compute_psds_from_scores(577 self.val_scores_postprocessed_buffer_student_strong,578 ground_truth,579 audio_durations,580 dtc_threshold=0.7,581 gtc_threshold=0.7,582 cttc_threshold=None,583 alpha_ct=0,584 alpha_st=1,585 # save_dir=os.path.join(save_dir, "", "scenario1"),586 )587 intersection_f1_macro_student = compute_per_intersection_macro_f1(588 self.val_buffer_student_strong,589 self.hparams["data"]["synth_val_tsv"],590 self.hparams["data"]["synth_val_dur"],591 )592 sed_eval_metrics_student = log_sedeval_metrics(593 self.val_buffer_student_strong[0.5],594 self.hparams["data"]["synth_val_tsv"],595 )596 strong_event_macro_student = sed_eval_metrics_student[0]597 strong_segment_macro_student = sed_eval_metrics_student[2]598 599 intersection_f1_macro_teacher = compute_per_intersection_macro_f1(600 self.val_buffer_teacher_strong,601 self.hparams["data"]["synth_val_tsv"],602 self.hparams["data"]["synth_val_dur"],603 )604 605 sed_eval_metrics_teacher = log_sedeval_metrics(606 self.val_buffer_teacher_strong[0.5],607 self.hparams["data"]["synth_val_tsv"],608 )609 strong_event_macro_teacher = sed_eval_metrics_teacher[0]610 strong_segment_macro_teacher = sed_eval_metrics_teacher[2]611 612 obj_metric_strong_type = self.hparams["training"].get("obj_metric_strong_type")613 if obj_metric_strong_type is None:614 strong_metric = psds1_sed_scores_eval_student615 elif obj_metric_strong_type == "event":616 strong_metric = strong_event_macro_student617 elif obj_metric_strong_type == "intersection":618 strong_metric = intersection_f1_macro_student619 elif obj_metric_strong_type == "psds":620 strong_metric = psds1_sed_scores_eval_student621 else:622 raise NotImplementedError(623 f"obj_metric_strong_type: {obj_metric_strong_type} not implemented."624 )625 626 obj_metric = torch.tensor(weak_student_f1_macro.item() + strong_metric)627 628 self.log("val/obj_metric", obj_metric, prog_bar=True, sync_dist=True)629 self.log("val/weak/student/macro_F1", weak_student_f1_macro, prog_bar=True)630 self.log("val/weak/teacher/macro_F1", weak_teacher_f1_macro)631 self.log(632 "val/strong/student/psds1_sed_scores_eval",633 psds1_sed_scores_eval_student,634 prog_bar=True,635 sync_dist=True,636 )637 self.log(638 "val/strong/student/intersection_f1_macro",639 intersection_f1_macro_student,640 prog_bar=True,641 sync_dist=True,642 )643 self.log(644 "val/strong/student/event_f1_macro",645 strong_event_macro_student,646 prog_bar=True,647 sync_dist=True,648 )649 self.log(650 "val/strong/student/segment_f1_macro",651 strong_segment_macro_student,652 prog_bar=True,653 sync_dist=True,654 )655 self.log(656 "val/strong/teacher/intersection_f1_macro",657 intersection_f1_macro_teacher,658 prog_bar=True,659 sync_dist=True,660 )661 self.log(662 "val/strong/teacher/event_f1_macro",663 strong_event_macro_teacher,664 prog_bar=True,665 sync_dist=True,666 )667 self.log(668 "val/strong/teacher/segment_f1_macro",669 strong_segment_macro_teacher,670 prog_bar=True,671 sync_dist=True,672 )673 674 # * free the buffers675 self.val_buffer_student_strong = {676 k: pd.DataFrame() for k in self.hparams["training"]["val_thresholds"]677 }678 self.val_buffer_teacher_strong = {679 k: pd.DataFrame() for k in self.hparams["training"]["val_thresholds"]680 }681 self.val_scores_postprocessed_buffer_student_strong = {}682 self.val_scores_postprocessed_buffer_teacher_strong = {}683 684 self.get_weak_student_f1_seg_macro.reset()685 self.get_weak_teacher_f1_seg_macro.reset()686 687 return obj_metric688 689 def on_save_checkpoint(self, checkpoint):690 checkpoint["sed_student"] = self.sed_student.state_dict()691 checkpoint["sed_teacher"] = self.sed_teacher.state_dict()692 return checkpoint693 694 def test_step(self, batch, batch_indx):695 """Apply Test to a batch (step), used only when (trainer.test is called)696 697 Args:698 batch: torch.Tensor, input batch tensor699 batch_indx: torch.Tensor, 1D tensor of indexes to know which data are present in each batch.700 Returns:701 """702 703 audio, labels, _, filenames = batch704 705 features = self.mel_spec(audio)706 preds_student, _ = self.sed_student(self.scaler(self.take_log(features)))707 preds_teacher, _ = self.sed_teacher(self.scaler(self.take_log(features)))708 709 if not self.evaluation:710 loss_student = self.supervised_loss(preds_student, labels)711 loss_teacher = self.supervised_loss(preds_teacher, labels)712 713 self.log("test/student/loss_strong", loss_student)714 self.log("test/teacher/loss_strong", loss_teacher)715 716 # * compute psds (Polyphonic Sound Detection Score)717 (718 scores_raw_student_strong,719 scores_postprocessed_student_strong,720 decoded_student_strong,721 ) = batched_decode_preds(722 preds_student,723 filenames,724 self.encoder,725 median_filter=self.hparams["training"]["median_window"],726 thresholds=list(self.test_psds_buffer_student.keys()) + [0.5],727 )728 729 self.test_scores_raw_buffer_student.update(scores_raw_student_strong)730 self.test_scores_postprocessed_buffer_student.update(731 scores_postprocessed_student_strong732 )733 for th in self.test_psds_buffer_student.keys():734 self.test_psds_buffer_student[th] = pd.concat(735 [self.test_psds_buffer_student[th], decoded_student_strong[th]],736 ignore_index=True,737 )738 739 (740 scores_raw_teacher_strong,741 scores_postprocessed_teacher_strong,742 decoded_teacher_strong,743 ) = batched_decode_preds(744 preds_teacher,745 filenames,746 self.encoder,747 median_filter=self.hparams["training"]["median_window"],748 thresholds=list(self.test_psds_buffer_teacher.keys()) + [0.5],749 )750 751 self.test_scores_raw_buffer_teacher.update(scores_raw_teacher_strong)752 self.test_scores_postprocessed_buffer_teacher.update(753 scores_postprocessed_teacher_strong754 )755 for th in self.test_psds_buffer_teacher.keys():756 self.test_psds_buffer_teacher[th] = pd.concat(757 [self.test_psds_buffer_teacher[th], decoded_teacher_strong[th]],758 ignore_index=True,759 )760 761 # compute f1 score762 self.decoded_student_05_buffer = pd.concat(763 [self.decoded_student_05_buffer, decoded_student_strong[0.5]]764 )765 self.decoded_teacher_05_buffer = pd.concat(766 [self.decoded_teacher_05_buffer, decoded_teacher_strong[0.5]]767 )768 769 def on_test_epoch_end(self):770 save_dir = os.path.join(self.exp_dir, "metrics_test")771 772 # * if evaluation is True, we only save the scores773 if self.evaluation:774 save_dir_raw = os.path.join(save_dir, "_scores", "raw")775 sed_scores_eval.io.write_sed_scores(776 self.test_scores_raw_buffer, save_dir_raw777 )778 print(f"\nRaw scores for saved in: {save_dir_raw}")779 780 save_dir_postprocessed = os.path.join(save_dir, "_scores", "postprocessed")781 sed_scores_eval.io.write_sed_scores(782 self.test_scores_postprocessed_buffer, save_dir_postprocessed783 )784 print(f"\nPostprocessed scores for saved in: {save_dir_postprocessed}")785 else:786 # * calculate the metrics and save them787 ground_truth = sed_scores_eval.io.read_ground_truth_events(788 self.hparams["data"]["test_tsv"]789 )790 audio_durations = sed_scores_eval.io.read_audio_durations(791 self.hparams["data"]["test_dur"]792 )793 if self.fast_dev_run:794 ground_truth = {795 audio_id: ground_truth[audio_id]796 for audio_id in self.test_scores_postprocessed_buffer_student797 }798 audio_durations = {799 audio_id: audio_durations[audio_id]800 for audio_id in self.test_scores_postprocessed_buffer_student801 }802 else:803 # drop audios without events804 ground_truth = {805 audio_id: gt for audio_id, gt in ground_truth.items() if len(gt) > 0806 }807 audio_durations = {808 audio_id: audio_durations[audio_id]809 for audio_id in ground_truth.keys()810 }811 psds1_student_psds_eval = compute_psds_from_operating_points(812 self.test_psds_buffer_student,813 self.hparams["data"]["test_tsv"],814 self.hparams["data"]["test_dur"],815 dtc_threshold=0.7,816 gtc_threshold=0.7,817 alpha_ct=0,818 alpha_st=1,819 save_dir=os.path.join(save_dir, "student", "scenario1"),820 )821 psds1_student_sed_scores_eval = compute_psds_from_scores(822 self.test_scores_postprocessed_buffer_student,823 ground_truth,824 audio_durations,825 dtc_threshold=0.7,826 gtc_threshold=0.7,827 cttc_threshold=None,828 alpha_ct=0,829 alpha_st=1,830 save_dir=os.path.join(save_dir, "student", "scenario1"),831 )832 psds2_student_psds_eval = compute_psds_from_operating_points(833 self.test_psds_buffer_student,834 self.hparams["data"]["test_tsv"],835 self.hparams["data"]["test_dur"],836 dtc_threshold=0.1,837 gtc_threshold=0.1,838 cttc_threshold=0.3,839 alpha_ct=0.5,840 alpha_st=1,841 save_dir=os.path.join(save_dir, "student", "scenario2"),842 )843 psds2_student_sed_scores_eval = compute_psds_from_scores(844 self.test_scores_postprocessed_buffer_student,845 ground_truth,846 audio_durations,847 dtc_threshold=0.1,848 gtc_threshold=0.1,849 cttc_threshold=0.3,850 alpha_ct=0.5,851 alpha_st=1,852 save_dir=os.path.join(save_dir, "student", "scenario2"),853 )854 psds1_teacher_psds_eval = compute_psds_from_operating_points(855 self.test_psds_buffer_teacher,856 self.hparams["data"]["test_tsv"],857 self.hparams["data"]["test_dur"],858 dtc_threshold=0.7,859 gtc_threshold=0.7,860 alpha_ct=0,861 alpha_st=1,862 save_dir=os.path.join(save_dir, "teacher", "scenario1"),863 )864 psds1_teacher_sed_scores_eval = compute_psds_from_scores(865 self.test_scores_postprocessed_buffer_teacher,866 ground_truth,867 audio_durations,868 dtc_threshold=0.7,869 gtc_threshold=0.7,870 cttc_threshold=None,871 alpha_ct=0,872 alpha_st=1,873 save_dir=os.path.join(save_dir, "teacher", "scenario1"),874 )875 psds2_teacher_psds_eval = compute_psds_from_operating_points(876 self.test_psds_buffer_teacher,877 self.hparams["data"]["test_tsv"],878 self.hparams["data"]["test_dur"],879 dtc_threshold=0.1,880 gtc_threshold=0.1,881 cttc_threshold=0.3,882 alpha_ct=0.5,883 alpha_st=1,884 save_dir=os.path.join(save_dir, "teacher", "scenario2"),885 )886 psds2_teacher_sed_scores_eval = compute_psds_from_scores(887 self.test_scores_postprocessed_buffer_teacher,888 ground_truth,889 audio_durations,890 dtc_threshold=0.1,891 gtc_threshold=0.1,892 cttc_threshold=0.3,893 alpha_ct=0.5,894 alpha_st=1,895 save_dir=os.path.join(save_dir, "teacher", "scenario2"),896 )897 898 sed_eval_metrics_student = log_sedeval_metrics(899 self.decoded_student_05_buffer,900 self.hparams["data"]["test_tsv"],901 os.path.join(save_dir, ""),902 )903 event_macro_student = sed_eval_metrics_student[0]904 segment_macro_student = sed_eval_metrics_student[2]905 906 sed_eval_metrics_teacher = log_sedeval_metrics(907 self.decoded_teacher_05_buffer,908 self.hparams["data"]["test_tsv"],909 os.path.join(save_dir, ""),910 )911 event_macro_teacher = sed_eval_metrics_teacher[0]912 segment_macro_teacher = sed_eval_metrics_teacher[2]913 914 intersection_f1_macro_student = compute_per_intersection_macro_f1(915 {"0.5": self.decoded_student_05_buffer},916 self.hparams["data"]["test_tsv"],917 self.hparams["data"]["test_dur"],918 )919 920 intersection_f1_macro_teacher = compute_per_intersection_macro_f1(921 {"0.5": self.decoded_teacher_05_buffer},922 self.hparams["data"]["test_tsv"],923 self.hparams["data"]["test_dur"],924 )925 926 results = {927 "test/student/psds1_psds_eval": psds1_student_psds_eval,928 "test/student/psds1_sed_scores_eval": psds1_student_sed_scores_eval,929 "test/student/psds2_psds_eval": psds2_student_psds_eval,930 "test/student/psds2_sed_scores_eval": psds2_student_sed_scores_eval,931 "test/student/segment_f1_macro": segment_macro_student,932 "test/student/event_f1_macro": event_macro_student,933 "test/student/intersection_f1_macro": intersection_f1_macro_student,934 "test/teacher/psds1_psds_eval": psds1_teacher_psds_eval,935 "test/teacher/psds1_sed_scores_eval": psds1_teacher_sed_scores_eval,936 "test/teacher/psds2_psds_eval": psds2_teacher_psds_eval,937 "test/teacher/psds2_sed_scores_eval": psds2_teacher_sed_scores_eval,938 "test/teacher/segment_f1_macro": segment_macro_teacher,939 "test/teacher/event_f1_macro": event_macro_teacher,940 "test/teacher/intersection_f1_macro": intersection_f1_macro_teacher,941 }942 943 if self.logger is not None:944 self.logger.log_metrics(results)945 self.logger.log_hyperparams(self.hparams, results)946 947 for key in results.keys():948 self.log(key, results[key], prog_bar=True, logger=True, sync_dist=True)949 950 def configure_optimizers(self):951 return [self.opt], [self.scheduler]952 953 def train_dataloader(self):954 self.train_loader = torch.utils.data.DataLoader(955 self.train_data,956 batch_sampler=self.train_sampler,957 num_workers=self.num_workers,958 )959 return self.train_loader960 961 def val_dataloader(self):962 self.val_loader = torch.utils.data.DataLoader(963 self.valid_data,964 batch_size=self.hparams["training"]["batch_size_val"],965 num_workers=self.num_workers,966 shuffle=False,967 drop_last=False,968 )969 return self.val_loader970 971 def test_dataloader(self):972 self.test_loader = torch.utils.data.DataLoader(973 self.test_data,974 batch_size=self.hparams["training"]["batch_size_val"],975 num_workers=self.num_workers,976 shuffle=False,977 drop_last=False,978 )979 return self.test_loader980 981 def forward(self, x, model="student"):982 features = self.mel_spec(x)983 features = features.unsqueeze(0)984 if model == "student":985 preds, _ = self.sed_student(self.scaler(self.take_log(features)))986 elif model == "teacher":987 preds, _ = self.sed_teacher(self.scaler(self.take_log(features)))988 else:989 raise NotImplementedError("Only student and teacher models are available")990 return preds991 992 993if __name__ == "__main__":994 with open("params.yaml", "r") as f:995 config = yaml.safe_load(f)996 997 encoder = ManyHotEncoder(998 list(classes_labels.keys()),999 audio_len=config["data"]["audio_max_len"],1000 frame_len=config["feats"]["n_filters"],1001 frame_hop=config["feats"]["hop_length"],1002 net_pooling=config["data"]["net_subsample"],1003 fs=config["data"]["fs"],1004 )1005 1006 sed = SED(config, encoder=encoder, sed=CRNN(**config["net"]))1007 print(sed.state_dict().keys())1008