zimoqingfeng/UMOP

error when num_classes not equal to 80

Closed this issue · 3 comments

File "/UMOP-master/mmdet/models/dense_heads/level_resampling_head.py", line 97, in loss_single
    pos_xse = - torch.log(prob) * y * label_weights[:, None]
RuntimeError: The size of tensor a (5) must match the size of tensor b (6) at non-singleton dimension 1

i set num_classes = 5 for my custom dataset, you can reproduce use

num_classes = 5 


classes=['person', 'bicycle', 'car', 'motorcycle', 'airplane'],
type=dataset_type,
ann_file='data/coco/annotations/instances_train2017.json',
img_prefix='data/coco/train2017/',

should modify here level_resampling_head.py#L94

y = F.one_hot(labels)[:, :80]

okey modify it as

y = F.one_hot(labels)[:, :self.num_classes]

Thx, the code has been modified for num_classes not equal to 80