|
import torch |
|
import torch.distributed.nn |
|
from torch import distributed as dist, nn as nn |
|
from torch.nn import functional as F |
|
import numpy as np |
|
from sklearn.metrics import average_precision_score, roc_auc_score, accuracy_score |
|
|
|
try: |
|
import horovod.torch as hvd |
|
except ImportError: |
|
hvd = None |
|
|
|
|
|
def gather_features( |
|
audio_features, |
|
text_features, |
|
audio_features_mlp=None, |
|
text_features_mlp=None, |
|
local_loss=False, |
|
gather_with_grad=False, |
|
rank=0, |
|
world_size=1, |
|
use_horovod=False, |
|
mlp_loss=False, |
|
): |
|
if use_horovod: |
|
assert hvd is not None, "Please install horovod" |
|
if gather_with_grad: |
|
all_audio_features = hvd.allgather(audio_features) |
|
all_text_features = hvd.allgather(text_features) |
|
if mlp_loss: |
|
all_audio_features_mlp = hvd.allgather(audio_features_mlp) |
|
all_text_features_mlp = hvd.allgather(text_features_mlp) |
|
else: |
|
with torch.no_grad(): |
|
all_audio_features = hvd.allgather(audio_features) |
|
all_text_features = hvd.allgather(text_features) |
|
if mlp_loss: |
|
all_audio_features_mlp = hvd.allgather(audio_features_mlp) |
|
all_text_features_mlp = hvd.allgather(text_features_mlp) |
|
if not local_loss: |
|
|
|
gathered_audio_features = list( |
|
all_audio_features.chunk(world_size, dim=0) |
|
) |
|
gathered_text_features = list( |
|
all_text_features.chunk(world_size, dim=0) |
|
) |
|
gathered_audio_features[rank] = audio_features |
|
gathered_text_features[rank] = text_features |
|
all_audio_features = torch.cat(gathered_audio_features, dim=0) |
|
all_text_features = torch.cat(gathered_text_features, dim=0) |
|
if mlp_loss: |
|
gathered_audio_features_mlp = list( |
|
all_audio_features_mlp.chunk(world_size, dim=0) |
|
) |
|
gathered_text_features_mlp = list( |
|
all_text_features_mlp.chunk(world_size, dim=0) |
|
) |
|
gathered_audio_features_mlp[rank] = audio_features_mlp |
|
gathered_text_features_mlp[rank] = text_features_mlp |
|
all_audio_features_mlp = torch.cat( |
|
gathered_audio_features_mlp, dim=0 |
|
) |
|
all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0) |
|
else: |
|
|
|
if gather_with_grad: |
|
all_audio_features = torch.cat( |
|
torch.distributed.nn.all_gather(audio_features), dim=0 |
|
) |
|
all_text_features = torch.cat( |
|
torch.distributed.nn.all_gather(text_features), dim=0 |
|
) |
|
if mlp_loss: |
|
all_audio_features_mlp = torch.cat( |
|
torch.distributed.nn.all_gather(audio_features_mlp), dim=0 |
|
) |
|
all_text_features_mlp = torch.cat( |
|
torch.distributed.nn.all_gather(text_features_mlp), dim=0 |
|
) |
|
else: |
|
gathered_audio_features = [ |
|
torch.zeros_like(audio_features) for _ in range(world_size) |
|
] |
|
gathered_text_features = [ |
|
torch.zeros_like(text_features) for _ in range(world_size) |
|
] |
|
dist.all_gather(gathered_audio_features, audio_features) |
|
dist.all_gather(gathered_text_features, text_features) |
|
if mlp_loss: |
|
gathered_audio_features_mlp = [ |
|
torch.zeros_like(audio_features_mlp) for _ in range(world_size) |
|
] |
|
gathered_text_features_mlp = [ |
|
torch.zeros_like(text_features_mlp) for _ in range(world_size) |
|
] |
|
dist.all_gather(gathered_audio_features_mlp, audio_features_mlp) |
|
dist.all_gather(gathered_text_features_mlp, text_features_mlp) |
|
if not local_loss: |
|
|
|
gathered_audio_features[rank] = audio_features |
|
gathered_text_features[rank] = text_features |
|
if mlp_loss: |
|
gathered_audio_features_mlp[rank] = audio_features_mlp |
|
gathered_text_features_mlp[rank] = text_features_mlp |
|
|
|
all_audio_features = torch.cat(gathered_audio_features, dim=0) |
|
all_text_features = torch.cat(gathered_text_features, dim=0) |
|
if mlp_loss: |
|
all_audio_features_mlp = torch.cat(gathered_audio_features_mlp, dim=0) |
|
all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0) |
|
if mlp_loss: |
|
return ( |
|
all_audio_features, |
|
all_text_features, |
|
all_audio_features_mlp, |
|
all_text_features_mlp, |
|
) |
|
else: |
|
return all_audio_features, all_text_features |
|
|
|
|
|
class ClipLoss(nn.Module): |
|
def __init__( |
|
self, |
|
local_loss=False, |
|
gather_with_grad=False, |
|
cache_labels=False, |
|
rank=0, |
|
world_size=1, |
|
use_horovod=False, |
|
mlp_loss=False, |
|
weight_loss_kappa=0, |
|
): |
|
super().__init__() |
|
self.local_loss = local_loss |
|
self.gather_with_grad = gather_with_grad |
|
self.cache_labels = cache_labels |
|
self.rank = rank |
|
self.world_size = world_size |
|
self.use_horovod = use_horovod |
|
self.mlp_loss = mlp_loss |
|
self.weighted_loss = bool(weight_loss_kappa != 0) |
|
self.weight_loss_kappa = weight_loss_kappa |
|
|
|
self.prev_num_logits = 0 |
|
self.labels = {} |
|
|
|
def forward( |
|
self, |
|
audio_features, |
|
text_features, |
|
logit_scale_a, |
|
logit_scale_t=None, |
|
audio_features_mlp=None, |
|
text_features_mlp=None, |
|
): |
|
device = audio_features.device |
|
if self.mlp_loss: |
|
if self.world_size > 1: |
|
( |
|
all_audio_features, |
|
all_text_features, |
|
all_audio_features_mlp, |
|
all_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=self.local_loss, |
|
gather_with_grad=self.gather_with_grad, |
|
rank=self.rank, |
|
world_size=self.world_size, |
|
use_horovod=self.use_horovod, |
|
mlp_loss=self.mlp_loss, |
|
) |
|
if self.local_loss: |
|
a_logits_per_audio = ( |
|
logit_scale_a * audio_features @ all_text_features_mlp.T |
|
) |
|
a_logits_per_text = ( |
|
logit_scale_a * text_features_mlp @ all_audio_features.T |
|
) |
|
t_logits_per_audio = ( |
|
logit_scale_t * audio_features_mlp @ all_text_features.T |
|
) |
|
t_logits_per_text = ( |
|
logit_scale_t * text_features @ all_audio_features_mlp.T |
|
) |
|
else: |
|
a_logits_per_audio = ( |
|
logit_scale_a * all_audio_features @ all_text_features_mlp.T |
|
) |
|
a_logits_per_text = a_logits_per_audio.T |
|
t_logits_per_audio = ( |
|
logit_scale_t * all_audio_features_mlp @ all_text_features.T |
|
) |
|
t_logits_per_text = t_logits_per_audio.T |
|
else: |
|
a_logits_per_audio = ( |
|
logit_scale_a * audio_features @ text_features_mlp.T |
|
) |
|
a_logits_per_text = logit_scale_a * text_features_mlp @ audio_features.T |
|
t_logits_per_audio = ( |
|
logit_scale_t * audio_features_mlp @ text_features.T |
|
) |
|
t_logits_per_text = logit_scale_t * text_features @ audio_features_mlp.T |
|
|
|
|
|
num_logits = a_logits_per_audio.shape[0] |
|
if self.prev_num_logits != num_logits or device not in self.labels: |
|
labels = torch.arange(num_logits, device=device, dtype=torch.long) |
|
if self.world_size > 1 and self.local_loss: |
|
labels = labels + num_logits * self.rank |
|
if self.cache_labels: |
|
self.labels[device] = labels |
|
self.prev_num_logits = num_logits |
|
else: |
|
labels = self.labels[device] |
|
|
|
if not self.weighted_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 |
|
else: |
|
audio_weight = (audio_features @ audio_features.T).detach() |
|
audio_weight = ( |
|
torch.exp( |
|
torch.sum(audio_weight, axis=1) |
|
/ (self.weight_loss_kappa * len(audio_weight)) |
|
) |
|
).detach() |
|
text_weight = (text_features @ text_features.T).detach() |
|
text_weight = ( |
|
torch.exp( |
|
torch.sum(text_weight, axis=1) |
|
/ (self.weight_loss_kappa * len(text_features)) |
|
) |
|
).detach() |
|
total_loss = ( |
|
F.cross_entropy(a_logits_per_audio, labels, weight=audio_weight) |
|
+ F.cross_entropy(a_logits_per_text, labels, weight=audio_weight) |
|
+ F.cross_entropy(t_logits_per_audio, labels, weight=text_weight) |
|
+ F.cross_entropy(t_logits_per_text, labels, weight=text_weight) |
|
) / 4 |
|
else: |
|
if self.world_size > 1: |
|
all_audio_features, all_text_features = gather_features( |
|
audio_features=audio_features, |
|
text_features=text_features, |
|
local_loss=self.local_loss, |
|
gather_with_grad=self.gather_with_grad, |
|
rank=self.rank, |
|
world_size=self.world_size, |
|
use_horovod=self.use_horovod, |
|
mlp_loss=self.mlp_loss, |
|
) |
|
|
|
if self.local_loss: |
|
logits_per_audio = ( |
|
logit_scale_a * audio_features @ all_text_features.T |
|
) |
|
logits_per_text = ( |
|
logit_scale_a * text_features @ all_audio_features.T |
|
) |
|
else: |
|
logits_per_audio = ( |
|
logit_scale_a * all_audio_features @ all_text_features.T |
|
) |
|
logits_per_text = logits_per_audio.T |
|
else: |
|
logits_per_audio = logit_scale_a * audio_features @ text_features.T |
|
logits_per_text = logit_scale_a * text_features @ audio_features.T |
|
|
|
|
|
num_logits = logits_per_audio.shape[0] |
|
if self.prev_num_logits != num_logits or device not in self.labels: |
|
labels = torch.arange(num_logits, device=device, dtype=torch.long) |
|
if self.world_size > 1 and self.local_loss: |
|
labels = labels + num_logits * self.rank |
|
if self.cache_labels: |
|
self.labels[device] = labels |
|
self.prev_num_logits = num_logits |
|
else: |
|
labels = self.labels[device] |
|
if not self.weighted_loss: |
|
total_loss = ( |
|
F.cross_entropy(logits_per_audio, labels) |
|
+ F.cross_entropy(logits_per_text, labels) |
|
) / 2 |
|
else: |
|
audio_weight = (all_audio_features @ all_audio_features.T).detach() |
|
audio_weight = ( |
|
torch.exp( |
|
torch.sum(audio_weight, axis=1) |
|
/ (self.weight_loss_kappa * len(all_audio_features)) |
|
) |
|
).detach() |
|
text_weight = (all_text_features @ all_text_features.T).detach() |
|
text_weight = ( |
|
torch.exp( |
|
torch.sum(text_weight, axis=1) |
|
/ (self.weight_loss_kappa * len(all_text_features)) |
|
) |
|
).detach() |
|
total_loss = ( |
|
F.cross_entropy(logits_per_audio, labels, weight=text_weight) |
|
+ F.cross_entropy(logits_per_text, labels, weight=audio_weight) |
|
) / 2 |
|
return total_loss |
|
|
|
|
|
def lp_gather_features(pred, target, world_size=1, use_horovod=False): |
|
if use_horovod: |
|
assert hvd is not None, "Please install horovod" |
|
with torch.no_grad(): |
|
all_preds = hvd.allgather(pred) |
|
all_targets = hvd.allgath(target) |
|
else: |
|
gathered_preds = [torch.zeros_like(pred) for _ in range(world_size)] |
|
gathered_targets = [torch.zeros_like(target) for _ in range(world_size)] |
|
|
|
dist.all_gather(gathered_preds, pred) |
|
dist.all_gather(gathered_targets, target) |
|
all_preds = torch.cat(gathered_preds, dim=0) |
|
all_targets = torch.cat(gathered_targets, dim=0) |
|
|
|
return all_preds, all_targets |
|
|
|
|
|
def get_map(pred, target): |
|
pred = torch.sigmoid(pred).numpy() |
|
target = target.numpy() |
|
return np.mean(average_precision_score(target, pred, average=None)) |
|
|
|
|
|
def get_acc(pred, target): |
|
pred = torch.argmax(pred, 1).numpy() |
|
target = torch.argmax(target, 1).numpy() |
|
return accuracy_score(target, pred) |
|
|
|
|
|
def get_mauc(pred, target): |
|
pred = torch.sigmoid(pred).numpy() |
|
target = target.numpy() |
|
return np.mean(roc_auc_score(target, pred, average=None)) |
|
|
|
|
|
class LPMetrics(object): |
|
def __init__(self, metric_names=["map", "acc", "mauc"]): |
|
self.metrics = [] |
|
for name in metric_names: |
|
self.metrics.append(self.get_metric(name)) |
|
self.metric_names = metric_names |
|
|
|
def get_metric(self, name): |
|
if name == "map": |
|
return get_map |
|
elif name == "acc": |
|
return get_acc |
|
elif name == "mauc": |
|
return get_mauc |
|
else: |
|
raise ValueError(f"the metric should be at least one of [map, acc, mauc]") |
|
|
|
def evaluate_mertics(self, pred, target): |
|
metric_dict = {} |
|
for i in range(len(self.metric_names)): |
|
metric_dict[self.metric_names[i]] = self.metrics[i](pred, target) |
|
return metric_dict |
|
|
|
|
|
def calc_celoss(pred, target): |
|
target = torch.argmax(target, 1).long() |
|
return nn.CrossEntropyLoss()(pred, target) |
|
|
|
|
|
class LPLoss(nn.Module): |
|
def __init__(self, loss_name): |
|
super().__init__() |
|
if loss_name == "bce": |
|
self.loss_func = nn.BCEWithLogitsLoss() |
|
elif loss_name == "ce": |
|
self.loss_func = calc_celoss |
|
elif loss_name == "mse": |
|
self.loss_func = nn.MSELoss() |
|
else: |
|
raise ValueError(f"the loss func should be at least one of [bce, ce, mse]") |
|
|
|
def forward(self, pred, target): |
|
loss = self.loss_func(pred, target) |
|
return loss |
|
|
