def hm_svoego_roda_loss(pred, target):
    pred_coord = heatmap_to_measure(pred)[0]
    target_coord = heatmap_to_measure(target)[0]

    pred = pred.relu() + 1e-15
    target[target < 1e-7] = 0
    target[target > 1 - 1e-7] = 1

    if torch.isnan(pred).any() or torch.isnan(target).any():
        return Loss.ZERO()

    bce = nn.BCELoss()(pred, target)

    if torch.isnan(bce).any():
        return Loss.ZERO()

    return Loss(bce + nn.MSELoss()(pred_coord, target_coord) * 0.0005)
def hm_svoego_roda_loss(pred, target, coef=1.0, l1_coef=0.0):
    pred_mes = UniformMeasure2DFactory.from_heatmap(pred)
    target_mes = UniformMeasure2DFactory.from_heatmap(target)

    # pred = pred.relu() + 1e-15
    # target[target < 1e-7] = 0
    # target[target > 1 - 1e-7] = 1

    if torch.isnan(pred).any() or torch.isnan(target).any():
        print("nan in hm")
        return Loss.ZERO()

    bce = nn.BCELoss()(pred, target)

    if torch.isnan(bce).any():
        print("nan in bce")
        return Loss.ZERO()

    return Loss(bce * coef + nn.MSELoss()(pred_mes.coord, target_mes.coord) *
                (0.0005 * coef) +
                nn.L1Loss()(pred_mes.coord, target_mes.coord) * l1_coef)
Esempio n. 3
0
 def add_generator_loss(self, loss: nn.Module, weight=1.0):
     return self.__add__(
         GANLossObject(
             lambda dx, dy: Loss.ZERO(), lambda dgz, real, fake: Loss(
                 loss(fake[0], real[0].detach()) * weight),
             self.discriminator))