mmpose-0.29.0/tests/test_losses/test_classification_loss.py

# Copyright (c) OpenMMLab. All rights reserved.
import torch


def test_bce_loss():
 from mmpose.models import build_loss

 # test BCE loss without target weight(None)
 loss_cfg = dict(type='BCELoss')
 loss = build_loss(loss_cfg)

 fake_pred = torch.zeros((1, 2))
 fake_label = torch.zeros((1, 2))
 assert torch.allclose(loss(fake_pred, fake_label), torch.tensor(0.))

 fake_pred = torch.ones((1, 2)) * 0.5
 fake_label = torch.zeros((1, 2))
 assert torch.allclose(
 loss(fake_pred, fake_label), -torch.log(torch.tensor(0.5)))

 # test BCE loss with target weight
 loss_cfg = dict(type='BCELoss', use_target_weight=True)
 loss = build_loss(loss_cfg)

 fake_pred = torch.ones((1, 2)) * 0.5
 fake_label = torch.zeros((1, 2))
 fake_weight = torch.ones((1, 2))
 assert torch.allclose(
 loss(fake_pred, fake_label, fake_weight),
 -torch.log(torch.tensor(0.5)))

 fake_weight[:, 0] = 0
 assert torch.allclose(
 loss(fake_pred, fake_label, fake_weight),
 -0.5 * torch.log(torch.tensor(0.5)))

 fake_weight = torch.ones(1)
 assert torch.allclose(
 loss(fake_pred, fake_label, fake_weight),
 -torch.log(torch.tensor(0.5)))