mirror of
https://github.com/JDAI-CV/fast-reid.git
synced 2025-06-03 14:50:47 +08:00
44cee30dfc
Summary: 1. refactor dataloader and heads 2. bugfix in fastattr, fastclas, fastface and partialreid 3. partial-fc supported in fastface
25 lines
712 B
Python
25 lines
712 B
Python
# encoding: utf-8
|
|
"""
|
|
@author: xingyu liao
|
|
@contact: sherlockliao01@gmail.com
|
|
"""
|
|
|
|
from fastreid.modeling.meta_arch import Baseline
|
|
from fastreid.modeling.meta_arch import META_ARCH_REGISTRY
|
|
|
|
|
|
@META_ARCH_REGISTRY.register()
|
|
class FaceBaseline(Baseline):
|
|
def __init__(self, cfg):
|
|
super().__init__(cfg)
|
|
self.pfc_enabled = cfg.MODEL.HEADS.PFC.ENABLED
|
|
|
|
def losses(self, outputs, gt_labels):
|
|
if not self.pfc_enabled:
|
|
return super().losses(outputs, gt_labels)
|
|
else:
|
|
# model parallel with partial-fc
|
|
# cls layer and loss computation in partial_fc.py
|
|
pred_features = outputs["features"]
|
|
return pred_features, gt_labels
|