import json
import logging
import math
import os
import time
from contextlib import suppress
import numpy as np
import torch
import torch.nn.functional as F
try:
import wandb
except ImportError:
wandb = None
from open_clip import ClipLoss, gather_features
from .distributed import is_master
from .zero_shot import zero_shot_eval
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def unwrap_model(model):
if hasattr(model, "module"):
return model.module
else:
return model
def train_one_epoch(
model, data, epoch, optimizer, scaler, scheduler, args, tb_writer=None
):
device = torch.device(args.device)
autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
model.train()
loss = ClipLoss(
local_loss=args.local_loss,
gather_with_grad=args.gather_with_grad,
cache_labels=True,
rank=args.rank,
world_size=args.world_size,
use_horovod=args.horovod,
mlp_loss=args.clap_mlploss,
weight_loss_kappa=args.kappa,
)
dataloader, sampler = data["train"].dataloader, data["train"].sampler
if args.distributed and sampler is not None:
sampler.set_epoch(epoch)
num_batches_per_epoch = dataloader.num_batches
sample_digits = math.ceil(math.log(dataloader.num_samples + 1, 10))
# for toy dataset
if args.dataset_type == "toy":
dataloader.dataset.generate_queue()
loss_m = AverageMeter()
batch_time_m = AverageMeter()
data_time_m = AverageMeter()
end = time.time()
for i, batch in enumerate(dataloader):
# logging.info(f"batch {i} of {num_batches_per_epoch}")
step = num_batches_per_epoch * epoch + i
if isinstance(scheduler, dict):
for s in scheduler.values():
s(step)
else:
scheduler(step)
audios = batch # contains mel_spec, wavform, and longer list
texts = batch["text"]
# audios = audios.to(device=device, non_blocking=True)
# texts = texts.to(device=device, non_blocking=True)
data_time_m.update(time.time() - end)
if isinstance(optimizer, dict):
for o_ in optimizer.values():
o_.zero_grad()
else:
optimizer.zero_grad()
with autocast():
(
audio_features,
text_features,
audio_features_mlp,
text_features_mlp,
logit_scale_a,
logit_scale_t,
) = model(audios, texts, device)
if args.clap_mlploss:
total_loss = loss(
audio_features=audio_features,
text_features=text_features,
logit_scale_a=logit_scale_a,
logit_scale_t=logit_scale_t,
audio_features_mlp=audio_features_mlp,
text_features_mlp=text_features_mlp,
)
else:
total_loss = loss(
audio_features=audio_features,
text_features=text_features,
logit_scale_a=logit_scale_a,
)
if isinstance(optimizer, dict):
if scaler is not None:
scaler.scale(total_loss).backward()
for o_ in optimizer.values():
if args.horovod:
o_.synchronize()
scaler.unscale_(o_)
with o_.skip_synchronize():
scaler.step(o_)
else:
scaler.step(o_)
scaler.update()
else:
total_loss.backward()
for o_ in optimizer.values():
o_.step()
else:
if scaler is not None:
scaler.scale(total_loss).backward()
if args.horovod:
optimizer.synchronize()
scaler.unscale_(optimizer)
with optimizer.skip_synchronize():
scaler.step(optimizer)
else:
scaler.step(optimizer)
scaler.update()
else:
total_loss.backward()
optimizer.step()
# Note: we clamp to 4.6052 = ln(100), as in the original paper.
with torch.no_grad():
unwrap_model(model).logit_scale_a.clamp_(0, math.log(100))
if args.clap_mlploss:
unwrap_model(model).logit_scale_t.clamp_(0, math.log(100))
batch_time_m.update(time.time() - end)
end = time.time()
batch_count = i + 1
if is_master(args) and (i % 100 == 0 or batch_count == num_batches_per_epoch):
if isinstance(audios, dict):
batch_size = len(audios["waveform"])
else:
batch_size = len(audios)
num_samples = batch_count * batch_size * args.world_size
samples_per_epoch = dataloader.num_samples
percent_complete = 100.0 * batch_count / num_batches_per_epoch
# NOTE loss is coarsely sampled, just master node and per log update
loss_m.update(total_loss.item(), batch_size)
logit_scale_scalar_a = logit_scale_a.item()
logit_scale_scalar_t = logit_scale_t.item()
if isinstance(optimizer, dict):
if args.clap_mlploss:
logging.info(
f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
f"Data (t): {data_time_m.avg:.3f} "
f"Batch (t): {batch_time_m.avg:.3f} "
f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]} "
f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
)
log_data = {
"loss": loss_m.val,
"data_time": data_time_m.val,
"batch_time": batch_time_m.val,
"scale_audio": logit_scale_scalar_a,
"scale_text": logit_scale_scalar_t,
"lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
}
else:
logging.info(
f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
f"Data (t): {data_time_m.avg:.3f} "
f"Batch (t): {batch_time_m.avg:.3f} "
f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]} "
f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
)
log_data = {
"loss": loss_m.val,
"data_time": data_time_m.val,
"batch_time": batch_time_m.val,
"scale_audio": logit_scale_scalar_a,
"lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
}
else:
if args.clap_mlploss:
logging.info(
f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
f"Data (t): {data_time_m.avg:.3f} "
f"Batch (t): {batch_time_m.avg:.3f} "
f"LR: {optimizer.param_groups[0]['lr']:5f} "
f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
)
# Save train loss / etc. Using non avg meter values as loggers have their own smoothing
log_data = {
"loss": loss_m.val,
"data_time": data_time_m.val,
"batch_time": batch_time_m.val,
"scale_audio": logit_scale_scalar_a,
"scale_text": logit_scale_scalar_t,
"lr": optimizer.param_groups[0]["lr"],
}
else:
logging.info(
f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
f"Data (t): {data_time_m.avg:.3f} "
f"Batch (t): {batch_time_m.avg:.3f} "
f"LR: {optimizer.param_groups[0]['lr']:5f} "
f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
)
# Save train loss / etc. Using non avg meter values as loggers have their own smoothing
log_data = {
"loss": loss_m.val,
"data_time": data_time_m.val,
"batch_time": batch_time_m.val,
"scale_audio": logit_scale_scalar_a,
"lr": optimizer.param_groups[0]["lr"],
}
for name, val in log_data.items():
name = "train/" + name
if tb_writer is not None:
tb_writer.add_scalar(name, val, step)
if args.wandb:
assert wandb is not None, "Please install wandb."
wandb.log({name: val, "step": step})
# resetting batch / data time meters per log window
batch_time_m.reset()
data_time_m.reset()
# end for
def evaluate(model, data, epoch, args, tb_writer=None):
metrics = {}
if not args.parallel_eval:
if not is_master(args):
return metrics
device = torch.device(args.device)
model.eval()
# CHANGE
# zero_shot_metrics = zero_shot_eval(model, data, epoch, args)
# metrics.update(zero_shot_metrics)
if is_master(args):
print("Evaluating...")
autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
if args.val_dataset_names == ["Clotho", "audiocaps"]:
# if only clotho and audiocaps are used, then we will use a different evaluation function.
# This is because in the Clotho and audiocaps valid and test set, there are 5 text for 1 audio.
if args.parallel_eval:
# (yusong): just a hack here. Don't use parallel eval when evaluating only clotho and audiocaps.
raise NotImplementedError(
"Parallel evaluation not supported for eval only Clotho and audiocaps."
)
val_metrics_per_dataset = evaluate_clotho_audiocaps(
model, data, epoch, args, autocast, device, tb_writer
)
for m in val_metrics_per_dataset.values():
metrics.update(m)
if "epoch" not in metrics.keys():
metrics.update({"epoch": epoch})
metrics = select_top_metric_clotho_audiocaps(
metrics, val_metrics_per_dataset, args
)
elif "val" in data and (
args.val_frequency
and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)
):
dataloader = data["val"].dataloader
num_samples = 0
samples_per_val = dataloader.num_samples
# FIXME this does not scale past small eval datasets
# all_audio_features @ all_text_features will blow up memory and compute very quickly
eval_info = {}
if args.clap_mlploss:
eval_info["all"] = {
"cumulative_loss": 0.0,
"num_samples": 0,
"all_audio_features": [],
"all_text_features": [],
"all_audio_features_mlp": [],
"all_text_features_mlp": [],
} # cumulative_loss = 0.0
else:
eval_info["all"] = {
"cumulative_loss": 0.0,
"num_samples": 0,
"all_audio_features": [],
"all_text_features": [],
} # cumu
# all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp = [], [], [], []
with torch.no_grad():
for i, batch in enumerate(dataloader):
audios = batch # contains mel_spec, wavform, and longer list
texts = batch["text"]
# audios = audios.to(device=device, non_blocking=True)
all_names = list(
set(["-".join(b.split("/")[-3:-1]) for b in batch["__url__"]])
)
for name in all_names:
if name not in eval_info.keys():
if args.clap_mlploss:
eval_info[name] = {
"cumulative_loss": 0.0,
"num_samples": 0,
"all_audio_features": [],
"all_text_features": [],
"all_audio_features_mlp": [],
"all_text_features_mlp": [],
}
else:
eval_info[name] = {
"cumulative_loss": 0.0,
"num_samples": 0,
"all_audio_features": [],
"all_text_features": [],
}
with autocast():
(
audio_features,
text_features,
audio_features_mlp,
text_features_mlp,
logit_scale_a,
logit_scale_t,
) = model(audios, texts, device)
if args.parallel_eval:
# multi-GPU eval
if args.clap_mlploss:
(
audio_features,
text_features,
audio_features_mlp,
text_features_mlp,
) = gather_features(
audio_features=audio_features,
text_features=text_features,
audio_features_mlp=audio_features_mlp,
text_features_mlp=text_features_mlp,
local_loss=False,
gather_with_grad=False,
rank=args.rank,
world_size=args.world_size,
use_horovod=args.horovod,
mlp_loss=args.clap_mlploss,
)
else:
(audio_features, text_features,) = gather_features(
audio_features=audio_features,
text_features=text_features,
local_loss=False,
gather_with_grad=False,
rank=args.rank,
world_size=args.world_size,
use_horovod=args.horovod,
mlp_loss=args.clap_mlploss,
)
if is_master(args):
num_samples += audio_features.shape[0]
for n in [*all_names, "all"]:
if n == "all":
eval_info[n]["all_audio_features"].append(
audio_features.cpu()
)
eval_info[n]["all_text_features"].append(
text_features.cpu()
)
if args.clap_mlploss:
eval_info[n]["all_audio_features_mlp"].append(
audio_features_mlp.cpu()
)
eval_info[n]["all_text_features_mlp"].append(
text_features_mlp.cpu()
)
else:
idx = np.where(
np.array(
[
"-".join(b.split("/")[-3:-1])
for b in batch["__url__"]
]
)
== n
)[0]
eval_info[n]["all_audio_features"].append(
audio_features.cpu().index_select(
0, torch.tensor(idx).long()
)
)
eval_info[n]["all_text_features"].append(
text_features.cpu().index_select(
0, torch.tensor(idx).long()
)
)
if args.clap_mlploss:
eval_info[n]["all_audio_features_mlp"].append(
audio_features_mlp.cpu().index_select(
0, torch.tensor(idx).long()
)
)
eval_info[n]["all_text_features_mlp"].append(
text_features_mlp.cpu().index_select(
0, torch.tensor(idx).long()
)
)
# print(f'eval step {i}') # (yusong): for debug
# cumulative_loss += total_loss * batch_size
# num_samples += batch_size
if is_master(args) and (i % 100) == 0: # and i != 0:
logging.info(
f"Eval Epoch: {epoch} [{num_samples} / {samples_per_val}]"
)
if is_master(args):
val_metrics_per_dataset = {}
for n in eval_info.keys():
if args.clap_mlploss:
metrics_single_dataset = get_metrics(
audio_features=torch.cat(
eval_info[n]["all_audio_features"]
),
text_features=torch.cat(eval_info[n]["all_text_features"]),
logit_scale_a=logit_scale_a.cpu(),
audio_features_mlp=torch.cat(
eval_info[n]["all_audio_features_mlp"]
),
text_features_mlp=torch.cat(
eval_info[n]["all_text_features_mlp"]
),
logit_scale_t=logit_scale_t.cpu(),
mlp_loss=args.clap_mlploss,
)
else:
metrics_single_dataset = get_metrics(
audio_features=torch.cat(
eval_info[n]["all_audio_features"]
),
text_features=torch.cat(eval_info[n]["all_text_features"]),
logit_scale_a=logit_scale_a.cpu(),
mlp_loss=args.clap_mlploss,
)
val_metrics_per_dataset[n] = {
n + "/" + k: v for k, v in metrics_single_dataset.items()
}
metrics.update(val_metrics_per_dataset[n])
if "epoch" not in metrics.keys():
metrics.update({"epoch": epoch})
if is_master(args):
if not metrics:
return metrics
logging.info(
f"Eval Epoch: {epoch} "
+ "\n".join(
[
"\t".join([f"{k}: {round(v, 4):.4f}" for k, v in m.items()])
for m in val_metrics_per_dataset.values()
]
)
)
if args.save_logs:
for name, val in metrics.items():
if tb_writer is not None:
tb_writer.add_scalar(f"val/{name}", val, epoch)
with open(os.path.join(args.checkpoint_path, "results.jsonl"), "a+") as f:
f.write(json.dumps(metrics))
f.write("\n")
if args.wandb:
assert wandb is not None, "Please install wandb."
for name, val in metrics.items():
wandb.log({f"val/{name}": val, "epoch": epoch})
return metrics
else:
return metrics
def get_metrics(
audio_features,
text_features,
logit_scale_a,
audio_features_mlp=None,
text_features_mlp=None,
logit_scale_t=None,
mlp_loss=False,
):
metrics = {}
if mlp_loss:
# Set up audio to text & text to audio similary matrice
a_logits_per_audio = (
(logit_scale_a * audio_features @ text_features_mlp.t()).detach().cpu()
)
a_logits_per_text = a_logits_per_audio.t().detach().cpu()
t_logits_per_audio = (
(logit_scale_t * audio_features_mlp @ text_features.t()).detach().cpu()
)
t_logits_per_text = t_logits_per_audio.t().detach().cpu()
labels = torch.arange(audio_features.shape[0]).long()
# Change the loss from two terms into four terms with 2x2 combined CE loss
total_loss = (
F.cross_entropy(a_logits_per_audio, labels)
+ F.cross_entropy(a_logits_per_text, labels)
+ F.cross_entropy(t_logits_per_audio, labels)
+ F.cross_entropy(t_logits_per_text, labels)
) / 4
metrics[f"cumulative_loss"] = total_loss.item()
metrics[f"num_samples"] = audio_features.shape[0]
logits = {
"audio_to_text": (a_logits_per_audio + t_logits_per_audio) / 2,
"text_to_audio": (a_logits_per_text + t_logits_per_text) / 2,
}
ground_truth = torch.arange(len(text_features)).view(-1, 1)
else:
# print("text_features", text_features)
# print("text_features.shape", text_features.shape)
logits_per_audio = (
(logit_scale_a * audio_features @ text_features.t()).detach().cpu()
)
logits_per_text = logits_per_audio.t().detach().cpu()
labels = torch.arange(audio_features.shape[0]).long()
# Change the loss from two terms into four terms with 2x2 combined CE loss
total_loss = (
F.cross_entropy(logits_per_audio, labels)
+ F.cross_entropy(logits_per_text, labels)
) / 2
metrics[f"cumulative_loss"] = total_loss.item()
metrics[f"num_samples"] = audio_features.shape[0]
logits = {"audio_to_text": logits_per_audio, "text_to_audio": logits_per_text}
ground_truth = torch.arange(len(text_features)).view(-1, 1)
for name, logit in logits.items():
ranking = torch.argsort(logit, descending=True)
preds = torch.where(ranking == ground_truth)[
1
] # (yusong) this line is slow because it uses single thread
preds = preds.detach().cpu().numpy()
metrics[f"{name}_mean_rank"] = preds.mean() + 1
metrics[f"{name}_median_rank"] = np.floor(np.median(preds)) + 1
for k in [1, 5, 10]:
metrics[f"{name}_R@{k}"] = np.mean(preds < k)
# map@10
metrics[f"{name}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
return metrics
def evaluate_clotho_audiocaps(
model, data, epoch, args, autocast, device, tb_writer=None
):
"""
Adapted from https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py.
1. for text-to-audio retrieval, do 5 times and average the results
2. for R@1, R@5, R@10 in audio-to-text retrieval, take the best rank among 5 text
3. for map@10 in audio-to-text retrieval:
3.1: sort the rank of 5 text
3.2: exclude the rank >=10 (0-index)
3.3: compute the map regarding the remaining ranks: np.mean(np.arange(1, len(ranks)+1) / ranks).
(3.3) That is, take the top ranks of 5 text that is < 10, and assign the descending number as ground truth.
(3.3) E.g.: the ground truth of first rank of the 5 text should be 1, the second rank should be 2, etc.
"""
# TODO: (yusong) only support single GPU evaluation and only support non-mlp case for now.
dataloader = data["val"].dataloader
with torch.no_grad():
eval_info = {}
for i, batch in enumerate(dataloader):
audios = batch # contains mel_spec, wavform, and longer list
# each item in the list has 5 texts
if args.tmodel == "transformer":
from open_clip import tokenize
texts = [tokenize(t) for t in batch["full_text"]]
texts = torch.cat(texts)
else:
from .data import tokenizer
texts = [
tokenizer(t) for t in batch["full_text"]
] # 5 texts for each audio
texts = {
k: torch.cat([t[k] for t in texts]) for k in texts[0].keys()
} # 5 x batch
# audios = audios.to(device=device, non_blocking=True)
all_names = list(
set(["-".join(b.split("/")[-3:-1]) for b in batch["__url__"]])
)
for name in all_names:
if name not in eval_info.keys():
# we will not use mlp outputs even if args.clap_mlploss=True
eval_info[name] = {
"cumulative_loss": 0.0,
"num_samples": 0,
"all_audio_features": [],
"all_text_features": [],
}
with autocast():
audio_features = model(audios, None, device)
text_features = model(None, texts, device)
audio_features = F.normalize(audio_features, dim=-1)
text_features = F.normalize(text_features, dim=-1)
all_names = list(
set(["-".join(b.split("/")[-3:-1]) for b in batch["__url__"]])
)
for n in all_names:
idx = np.where(
np.array(
["-".join(b.split("/")[-3:-1]) for b in batch["__url__"]]
)
== n
)[0]
eval_info[n]["all_audio_features"].append(
audio_features.cpu().index_select(0, torch.tensor(idx).long())
)
# (yusong) please double-check. This is for selecting 5 text features at once.
# because idx is a list of indices in size of num_samples,
# and text_features is a tensor of size (5*num_samples, dim)
# so we need to select 5 consecutive indices at once for a single index in idx.
eval_info[n]["all_text_features"].append(
text_features.cpu()
.reshape([-1, 5, text_features.shape[1]])
.index_select(0, torch.tensor(idx).long())
.reshape([-1, text_features.shape[1]])
)
val_metrics_all = {}
for n in eval_info.keys():
logit_scale_a, logit_scale_t = model(None, None, device)
logit_scale_a = logit_scale_a.cpu()
audio_features = torch.cat(eval_info[n]["all_audio_features"], dim=0)
text_features = torch.cat(eval_info[n]["all_text_features"], dim=0)
logits_per_audio = (
(logit_scale_a * audio_features @ text_features.t()).detach().cpu()
)
logits_per_text = logits_per_audio.t().detach().cpu()
# logits_per_audio shape: [num_samples, num_samples*5]
# logits_per_text shape: [num_samples*5, num_samples]
logging.info(
f"dataset {n}, logits_per_audio shape: {logits_per_audio.shape}, "
f"logits_per_text shape: {logits_per_text.shape}"
)
metrics = {}
num_samples = audio_features.shape[0]
metrics[f"num_samples"] = num_samples
# (yusong) the following code is very important, please double-check:
# logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d]
# logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
# Those two are retrieving one of the 5 text for each audio.
labels = torch.arange(audio_features.shape[0]).long()
audio_to_text_loss = [
F.cross_entropy(
logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d],
labels,
)
for d in range(5)
]
text_to_audio_loss = [
F.cross_entropy(
logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :],
labels,
)
for d in range(5)
]
total_loss = (np.mean(audio_to_text_loss) + np.mean(text_to_audio_loss)) / 2
metrics[f"cumulative_loss"] = total_loss.item()
# text to audio: do 5 times
pred_text = []
for d in range(5):
logit = logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
ground_truth = torch.arange(len(logit)).view(-1, 1)
ranking = torch.argsort(
logit, descending=True
) # [num_samples, num_samples]
preds = torch.where(ranking == ground_truth)[1]
pred_text.append(preds.detach().cpu().numpy())
pred_text_concat = np.concatenate(pred_text, axis=0) # [5*num_samples]
metrics[f"text_to_audio_mean_rank"] = pred_text_concat.mean() + 1
metrics[f"text_to_audio_median_rank"] = (
np.floor(np.median(pred_text_concat)) + 1
)
for k in [1, 5, 10]:
metrics[f"text_to_audio_R@{k}"] = np.mean(pred_text_concat < k)
# map@10
metrics[f"text_to_audio_mAP@10"] = np.mean(
np.where(pred_text_concat < 10, 1 / (pred_text_concat + 1), 0.0)
)
# audio to text: take the best result
# for audio to text map 10, sort and assign descending ground truth.
# see https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py#L103
# map@10
map_all = []
pred_audio_all = []
for d in range(num_samples):
# logits_per_audio: [num_samples, num_samples*5]
logit_single = logits_per_audio[d, :] # [5*num_samples]
# Ground-truth index: [d*5, d*5+1, d*5+2, d*5+3, d*5+4]
ranking = torch.argsort(
logit_single, descending=True
) # [5*num_samples]
# ranking: the index of first match, second match, ...
ground_truth = torch.arange(d * 5, d * 5 + 5)[None]
all_pred = torch.where(
torch.stack([ranking] * 5) == ground_truth.view(-1, 1)
)[1]
min_pred = torch.min(all_pred)
pred_audio_all.append(min_pred.detach().cpu().numpy())
all_pred_filter = all_pred[all_pred < 10].detach().cpu().numpy()
# /5 because we have 5 text, so it means for the text rank >=10 we count as 0.
map_single = (
np.sum(
(np.arange(1, len(all_pred_filter) + 1) / (all_pred_filter + 1))
)
/ 5
)
map_all.append(map_single)
metrics[f"audio_to_text_mAP@10"] = np.mean(map_all)
for k in [1, 5, 10]:
metrics[f"audio_to_text_R@{k}"] = np.mean(np.array(pred_audio_all) < k)
val_metrics_all[n] = {n + "/" + k: v for k, v in metrics.items()}
return val_metrics_all
def calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset):
"""
Calculate performance for Clotho+AudioCaps for model selection.
"""
selection_performance_all = []
for n in val_metrics_per_dataset.keys():
selection_performance = (
val_metrics_per_dataset[n][f"{n}/audio_to_text_mAP@10"]
+ val_metrics_per_dataset[n][f"{n}/text_to_audio_mAP@10"]
) / 2
selection_performance_all.append(selection_performance)
return np.mean(selection_performance_all)
def select_top_metric_clotho_audiocaps(metrics, val_metrics_per_dataset, args):
# val_metrics_per_dataset: dict, key: dataset name, value: dict, key: metric name, value: metric value
# metrics: dict, key: metric name, value: metric value
# Hack: use args to save the top performance
if not hasattr(args, "top_selection_performance"):
selection_performance = calculate_selection_performance_clotho_audiocaps(
val_metrics_per_dataset
)
# TODO: write the if and else together
metric_update = {}
for n in val_metrics_per_dataset.keys():
for k in val_metrics_per_dataset[n].keys():
metric_update[
k.split("/")[0] + "-top" + "/" + k.split("/")[1]
] = val_metrics_per_dataset[n][k]
metric_update["top_selection_performance"] = selection_performance
metric_update["top-selection-epoch"] = metrics["epoch"]
metrics.update(metric_update)
args.top_metric = metric_update
args.top_selection_performance = selection_performance
else:
selection_performance_new = calculate_selection_performance_clotho_audiocaps(
val_metrics_per_dataset
)
selection_performance_old = args.top_selection_performance
if selection_performance_new > selection_performance_old:
metric_update = {}
for n in val_metrics_per_dataset.keys():
for k in val_metrics_per_dataset[n].keys():
metric_update[
k.split("/")[0] + "-top" + "/" + k.split("/")[1]
] = val_metrics_per_dataset[n][k]
metric_update["top_selection_performance"] = selection_performance_new
metric_update["top-selection-epoch"] = metrics["epoch"]
metrics.update(metric_update)
args.top_metric = metric_update
args.top_selection_performance = selection_performance_new
else:
metrics.update(args.top_metric)
return metrics