import torch import torch.distributed as dist import torch.nn.functional as F from lightning.pytorch import LightningModule from torchmetrics.classification import ( BinaryConfusionMatrix, BinaryF1Score, BinaryPrecision, BinaryRecall, ) from models.asd import AudioVisualASD from tasks.asd_metrics import ava_active_speaker_map from tasks.ava_evaluation import export_ava_predictions from tasks.setup import ( get_adamw_optimizer, get_cosine_warmup_scheduler, get_optimizer_cfg, get_param_groups, ) class ASDTask(LightningModule): def __init__(self, cfg): super().__init__() self.save_hyperparameters(cfg) self.model = AudioVisualASD(cfg["asd"]) self.precision = BinaryPrecision() self.recall = BinaryRecall() self.f1 = BinaryF1Score() self.confusion = BinaryConfusionMatrix() self.validation_entities = [] evaluation = cfg.get("evaluation", {}) self.ava_groundtruth = evaluation.get("ava_groundtruth") self.ava_predictions = evaluation.get("ava_predictions") def forward(self, audio, visual): return self.model(audio, visual, return_embeddings=True) @staticmethod def _masked_bce(logits, targets, scored): """Average binary cross-entropy over scored frames only.""" loss = F.binary_cross_entropy_with_logits( logits, targets.float(), reduction="none" ) weight = scored.float() if loss.dim() == 3: weight = weight.unsqueeze(-1) weight = weight.expand_as(loss) return (loss * weight).sum() / weight.sum().clamp_min(1.0) def _shared_step(self, batch): audio, visual, projection_labels, activity_labels, scored = batch[:5] projection, audio_logits, visual_logits, _, audio_emb, visual_emb = self( audio, visual ) audio_logits = audio_logits.squeeze(-1) visual_logits = visual_logits.squeeze(-1) activity_logits = (audio_logits + visual_logits) / 2 projection_loss = self._masked_bce(projection, projection_labels, scored) audio_loss = self._masked_bce(audio_logits, activity_labels, scored) visual_loss = self._masked_bce(visual_logits, activity_labels, scored) contrastive_loss = self._talknce( audio_emb, visual_emb, activity_labels, scored ) loss = projection_loss + 0.4 * audio_loss + 0.4 * visual_loss + 0.3 * contrastive_loss losses = { "loss": loss, "projection_loss": projection_loss, "audio_loss": audio_loss, "visual_loss": visual_loss, "contrastive_loss": contrastive_loss, } return activity_logits, losses def training_step(self, batch, batch_idx): _, losses = self._shared_step(batch) for name, value in losses.items(): self.log(f"train/{name}", value, prog_bar=name == "loss", sync_dist=True) return losses["loss"] def validation_step(self, batch, batch_idx): logits, losses = self._shared_step(batch) labels = batch[3].int() probabilities = logits.sigmoid() if len(batch) == 6: scores = probabilities.detach().cpu() targets = labels.detach().cpu() for index, metadata in enumerate(batch[5]): frames = min(int(metadata["model_frames"]), scores.shape[-1]) self.validation_entities.append( ( metadata["sample_id"], scores[index, :frames], targets[index, :frames], float(metadata["source_fps"]), ) ) scored = batch[4].bool() flat_scores, flat_labels = probabilities[scored], labels[scored] self.precision.update(flat_scores, flat_labels) self.recall.update(flat_scores, flat_labels) self.f1.update(flat_scores, flat_labels) self.confusion.update(flat_scores, flat_labels) for name, value in losses.items(): self.log(f"val/{name}", value, prog_bar=name == "loss", sync_dist=True) def _validation_is_partial(self): """True when this pass deliberately saw only part of the val set. The official export needs one score for every groundtruth entity, so it cannot run during the sanity check or a truncated debug pass. An unexpectedly incomplete set still raises inside the exporter. """ return bool( self.trainer.sanity_checking or self.trainer.fast_dev_run or self.trainer.limit_val_batches != 1.0 ) def on_validation_epoch_end(self): if dist.is_available() and dist.is_initialized(): gathered = [None] * dist.get_world_size() dist.all_gather_object(gathered, self.validation_entities) entities = [entity for item in gathered for entity in item] else: entities = self.validation_entities # The distributed val sampler pads the epoch by repeating samples, so # the same entity can arrive from several ranks. Ranks are not required # to agree bit for bit, so the first copy wins. scores_by_entity = {} labels_by_entity = {} fps_by_entity = {} for entity_id, entity_scores, entity_labels, source_fps in entities: if entity_id in scores_by_entity: continue scores_by_entity[entity_id] = entity_scores labels_by_entity[entity_id] = entity_labels fps_by_entity[entity_id] = source_fps if scores_by_entity: scores = torch.cat(list(scores_by_entity.values())) labels = torch.cat(list(labels_by_entity.values())) else: scores = torch.empty(0) labels = torch.empty(0, dtype=torch.int) map_25hz = ava_active_speaker_map(scores, labels) matrix = self.confusion.compute() false_positive_rate = matrix[0, 1] / matrix[0].sum().clamp_min(1) self.log("val/mAP_25hz", map_25hz, prog_bar=True) if ( self.ava_groundtruth and self.ava_predictions and scores_by_entity and not self._validation_is_partial() ): # Only rank zero writes the prediction file, but every rank has to # log the metric or a ModelCheckpoint monitoring it will hang. official_map = None if self.trainer.is_global_zero: official_map = export_ava_predictions( self.ava_groundtruth, self.ava_predictions, scores_by_entity, fps_by_entity=fps_by_entity, ) if dist.is_available() and dist.is_initialized(): payload = [official_map] dist.broadcast_object_list(payload, src=0) official_map = payload[0] self.log("val/mAP_official", official_map, prog_bar=True) self.log("val/precision", self.precision.compute(), sync_dist=True) self.log("val/recall", self.recall.compute(), sync_dist=True) self.log("val/f1", self.f1.compute(), prog_bar=True, sync_dist=True) self.log("val/fpr", false_positive_rate, sync_dist=True) for metric in ( self.precision, self.recall, self.f1, self.confusion, ): metric.reset() self.validation_entities.clear() @staticmethod def _talknce(audio, visual, labels, scored=None, temperature=0.07): positive = labels.bool() if scored is not None: positive = positive & scored.bool() if audio.shape[0] < 2 or not positive.any(): return audio.sum() * 0 audio = F.normalize(audio, dim=-1) visual = F.normalize(visual, dim=-1) logits = torch.einsum("btd,ctd->tbc", audio, visual) / temperature batch_size = audio.shape[0] targets = torch.arange(batch_size, device=audio.device).repeat(audio.shape[1]) active = positive.transpose(0, 1).reshape(-1) return F.cross_entropy(logits.reshape(-1, batch_size)[active], targets[active]) def configure_optimizers(self): cfg = self.hparams["asd"] weight_decay = float(cfg["weight_decay"]) optimizer = get_adamw_optimizer( get_param_groups(self.named_parameters(), weight_decay), float(cfg["lr"]), weight_decay, ) scheduler = get_cosine_warmup_scheduler( optimizer, self.trainer.estimated_stepping_batches, float(cfg["warmup_ratio"]), float(cfg["min_lr"]) if cfg.get("min_lr") else None, ) return get_optimizer_cfg(optimizer, scheduler)