radames/Text2Human-API

1import torch.nn as nn2import torch.nn.functional as F3 4 5class BCELoss(nn.Module):6 7    def forward(self, prediction, target):8        loss = F.binary_cross_entropy_with_logits(prediction, target)9        return loss, {}10 11 12class BCELossWithQuant(nn.Module):13 14    def __init__(self, codebook_weight=1.):15        super().__init__()16        self.codebook_weight = codebook_weight17 18    def forward(self, qloss, target, prediction, split):19        bce_loss = F.binary_cross_entropy_with_logits(prediction, target)20        loss = bce_loss + self.codebook_weight * qloss21        return loss, {22            "{}/total_loss".format(split): loss.clone().detach().mean(),23            "{}/bce_loss".format(split): bce_loss.detach().mean(),24            "{}/quant_loss".format(split): qloss.detach().mean()25        }26