mirror of
https://github.com/open-mmlab/mmclassification.git
synced 2025-06-03 21:53:55 +08:00
07bb15e5fd
* resolve conflicts add heads and config for multilabel tasks * minor change * remove evaluating mAP in head * add baseline config * add configs * reserve only one config * minor change * fix minor bug * minor change * minor change * add unittests and fix docstrings
23 lines
607 B
Python
23 lines
607 B
Python
import torch
|
|
|
|
from mmcls.models.heads import MultiLabelClsHead, MultiLabelLinearClsHead
|
|
|
|
|
|
def test_multilabel_head():
|
|
head = MultiLabelClsHead()
|
|
fake_cls_score = torch.rand(4, 3)
|
|
fake_gt_label = torch.randint(0, 2, (4, 3))
|
|
|
|
losses = head.loss(fake_cls_score, fake_gt_label)
|
|
assert losses['loss'].item() > 0
|
|
|
|
|
|
def test_multilabel_linear_head():
|
|
head = MultiLabelLinearClsHead(3, 5)
|
|
fake_cls_score = torch.rand(4, 3)
|
|
fake_gt_label = torch.randint(0, 2, (4, 3))
|
|
|
|
head.init_weights()
|
|
losses = head.loss(fake_cls_score, fake_gt_label)
|
|
assert losses['loss'].item() > 0
|