mirrors
/
fast-reid
mirror of https://github.com/JDAI-CV/fast-reid.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix($modeling/heads/bnneck_head): fix heads outputs bug 

fix bug of heads outputs, which will lead to no targets return.
pull/44/head
liaoxingyu 2020-04-27 11:48:21 +08:00
parent a6bd0371e2
commit 2efbc6d371
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
fastreid/modeling/heads/bnneck_head.py
View File

@ -22,6 +22,7 @@ class BNneckHead(nn.Module):
self.bnneck = NoBiasBatchNorm1d(in_feat)
self.bnneck.apply(weights_init_kaiming)
# identity classification layer
if cfg.MODEL.HEADS.CLS_LAYER == 'linear':
self.classifier = nn.Linear(in_feat, self._num_classes, bias=False)
elif cfg.MODEL.HEADS.CLS_LAYER == 'arcface':
@ -45,4 +46,4 @@ class BNneckHead(nn.Module):
pred_class_logits = self.classifier(bn_feat)
except TypeError:
pred_class_logits = self.classifier(bn_feat, targets)
return pred_class_logits, global_feat
return pred_class_logits, bn_feat, targets