From 44cee30dfc929df2051d1ca56e36dc152866e96b Mon Sep 17 00:00:00 2001 From: liaoxingyu Date: Fri, 2 Apr 2021 21:33:13 +0800 Subject: [PATCH] update fastreid v1.2 Summary: 1. refactor dataloader and heads 2. bugfix in fastattr, fastclas, fastface and partialreid 3. partial-fc supported in fastface --- configs/VeRi/sbs_R50-ibn.yml | 1 - fastreid/config/defaults.py | 4 +- fastreid/data/build.py | 41 ++-- fastreid/engine/defaults.py | 4 +- fastreid/engine/hooks.py | 1 + fastreid/layers/any_softmax.py | 113 ++++------ fastreid/modeling/heads/embedding_head.py | 10 +- projects/FastAttr/fastattr/__init__.py | 4 +- projects/FastAttr/fastattr/config.py | 3 +- projects/FastAttr/fastattr/data_build.py | 74 ------- .../FastAttr/fastattr/modeling/__init__.py | 9 + .../fastattr/{ => modeling}/attr_baseline.py | 0 .../fastattr/{ => modeling}/attr_head.py | 0 .../fastattr/{ => modeling}/bce_loss.py | 0 projects/FastAttr/train_net.py | 49 ++++- projects/FastClas/configs/base-clas.yaml | 23 +- projects/FastClas/fastclas/dataset.py | 2 +- projects/FastClas/train_net.py | 39 +++- projects/FastFace/README.md | 5 +- projects/FastFace/configs/face_base.yml | 41 ++-- projects/FastFace/configs/r101_ir.yml | 5 +- projects/FastFace/configs/r50_ir.yml | 7 +- projects/FastFace/fastface/__init__.py | 7 +- projects/FastFace/fastface/build.py | 46 ---- projects/FastFace/fastface/config.py | 16 ++ projects/FastFace/fastface/datasets/ms1mv2.py | 2 - projects/FastFace/fastface/face_data.py | 80 +++++++ .../FastFace/fastface/modeling/__init__.py | 10 + .../fastface/modeling/face_baseline.py | 24 +++ .../FastFace/fastface/modeling/face_head.py | 39 ++++ .../FastFace/fastface/modeling/partial_fc.py | 196 ++++++++++++++++++ .../fastface/{ => modeling}/resnet_ir.py | 0 projects/FastFace/fastface/trainer.py | 173 ++++++++++++++++ projects/FastFace/train_net.py | 31 +-- .../PartialReID/configs/partial_market.yml | 72 ++++--- projects/PartialReID/partialreid/config.py | 4 +- .../PartialReID/partialreid/dsr_distance.py | 4 +- .../PartialReID/partialreid/dsr_evaluation.py | 42 ++-- projects/PartialReID/partialreid/dsr_head.py | 98 ++++----- .../partialreid/partialbaseline.py | 61 ++---- 40 files changed, 862 insertions(+), 478 deletions(-) delete mode 100644 projects/FastAttr/fastattr/data_build.py create mode 100644 projects/FastAttr/fastattr/modeling/__init__.py rename projects/FastAttr/fastattr/{ => modeling}/attr_baseline.py (100%) rename projects/FastAttr/fastattr/{ => modeling}/attr_head.py (100%) rename projects/FastAttr/fastattr/{ => modeling}/bce_loss.py (100%) delete mode 100644 projects/FastFace/fastface/build.py create mode 100644 projects/FastFace/fastface/config.py create mode 100644 projects/FastFace/fastface/face_data.py create mode 100644 projects/FastFace/fastface/modeling/__init__.py create mode 100644 projects/FastFace/fastface/modeling/face_baseline.py create mode 100644 projects/FastFace/fastface/modeling/face_head.py create mode 100644 projects/FastFace/fastface/modeling/partial_fc.py rename projects/FastFace/fastface/{ => modeling}/resnet_ir.py (100%) create mode 100644 projects/FastFace/fastface/trainer.py diff --git a/configs/VeRi/sbs_R50-ibn.yml b/configs/VeRi/sbs_R50-ibn.yml index 71ea633c..967d67db 100644 --- a/configs/VeRi/sbs_R50-ibn.yml +++ b/configs/VeRi/sbs_R50-ibn.yml @@ -11,7 +11,6 @@ MODEL: SOLVER: OPT: SGD - NESTEROV: True BASE_LR: 0.01 ETA_MIN_LR: 7.7e-5 diff --git a/fastreid/config/defaults.py b/fastreid/config/defaults.py index d5ce1a4c..ae76a063 100644 --- a/fastreid/config/defaults.py +++ b/fastreid/config/defaults.py @@ -73,8 +73,8 @@ _C.MODEL.HEADS.CLS_LAYER = "Linear" # ArcSoftmax" or "CircleSoftmax" # Margin and Scale for margin-based classification layer -_C.MODEL.HEADS.MARGIN = 0.15 -_C.MODEL.HEADS.SCALE = 128 +_C.MODEL.HEADS.MARGIN = 0. +_C.MODEL.HEADS.SCALE = 1 # ---------------------------------------------------------------------------- # # REID LOSSES options diff --git a/fastreid/data/build.py b/fastreid/data/build.py index 51747087..4d508575 100644 --- a/fastreid/data/build.py +++ b/fastreid/data/build.py @@ -26,21 +26,19 @@ _root = os.getenv("FASTREID_DATASETS", "datasets") -def _train_loader_from_config(cfg, *, Dataset=None, transforms=None, sampler=None, **kwargs): +def _train_loader_from_config(cfg, *, train_set=None, transforms=None, sampler=None, **kwargs): if transforms is None: transforms = build_transforms(cfg, is_train=True) - if Dataset is None: - Dataset = CommDataset + if train_set is None: + train_items = list() + for d in cfg.DATASETS.NAMES: + data = DATASET_REGISTRY.get(d)(root=_root, **kwargs) + if comm.is_main_process(): + data.show_train() + train_items.extend(data.train) - train_items = list() - for d in cfg.DATASETS.NAMES: - data = DATASET_REGISTRY.get(d)(root=_root, **kwargs) - if comm.is_main_process(): - data.show_train() - train_items.extend(data.train) - - train_set = Dataset(train_items, transforms, relabel=True) + train_set = CommDataset(train_items, transforms, relabel=True) if sampler is None: sampler_name = cfg.DATALOADER.SAMPLER_TRAIN @@ -92,24 +90,25 @@ def build_reid_train_loader( return train_loader -def _test_loader_from_config(cfg, dataset_name, *, Dataset=None, transforms=None, **kwargs): +def _test_loader_from_config(cfg, *, dataset_name=None, test_set=None, num_query=0, transforms=None, **kwargs): if transforms is None: transforms = build_transforms(cfg, is_train=False) - if Dataset is None: - Dataset = CommDataset - - data = DATASET_REGISTRY.get(dataset_name)(root=_root, **kwargs) - if comm.is_main_process(): - data.show_test() - test_items = data.query + data.gallery + if test_set is None: + assert dataset_name is not None, "dataset_name must be explicitly passed in when test_set is not provided" + data = DATASET_REGISTRY.get(dataset_name)(root=_root, **kwargs) + if comm.is_main_process(): + data.show_test() + test_items = data.query + data.gallery + test_set = CommDataset(test_items, transforms, relabel=False) - test_set = Dataset(test_items, transforms, relabel=False) + # Update query number + num_query = len(data.query) return { "test_set": test_set, "test_batch_size": cfg.TEST.IMS_PER_BATCH, - "num_query": len(data.query), + "num_query": num_query, } diff --git a/fastreid/engine/defaults.py b/fastreid/engine/defaults.py index c0975d77..e11e8b01 100644 --- a/fastreid/engine/defaults.py +++ b/fastreid/engine/defaults.py @@ -85,7 +85,7 @@ def default_setup(cfg, args): PathManager.mkdirs(output_dir) rank = comm.get_rank() - setup_logger(output_dir, distributed_rank=rank, name="fvcore") + # setup_logger(output_dir, distributed_rank=rank, name="fvcore") logger = setup_logger(output_dir, distributed_rank=rank) logger.info("Rank of current process: {}. World size: {}".format(rank, comm.get_world_size())) @@ -423,7 +423,7 @@ def build_test_loader(cls, cfg, dataset_name): It now calls :func:`fastreid.data.build_reid_test_loader`. Overwrite it if you'd like a different data loader. """ - return build_reid_test_loader(cfg, dataset_name) + return build_reid_test_loader(cfg, dataset_name=dataset_name) @classmethod def build_evaluator(cls, cfg, dataset_name, output_dir=None): diff --git a/fastreid/engine/hooks.py b/fastreid/engine/hooks.py index 0ebb81b4..76b641bc 100644 --- a/fastreid/engine/hooks.py +++ b/fastreid/engine/hooks.py @@ -360,6 +360,7 @@ def _do_eval(self): ) self.trainer.storage.put_scalars(**flattened_results, smoothing_hint=False) + torch.cuda.empty_cache() # Evaluation may take different time among workers. # A barrier make them start the next iteration together. comm.synchronize() diff --git a/fastreid/layers/any_softmax.py b/fastreid/layers/any_softmax.py index 9b493581..9d643bda 100644 --- a/fastreid/layers/any_softmax.py +++ b/fastreid/layers/any_softmax.py @@ -4,71 +4,57 @@ @contact: sherlockliao01@gmail.com """ -import math - import torch import torch.nn as nn -import torch.nn.functional as F __all__ = [ - 'Linear', - 'ArcSoftmax', - 'CosSoftmax', - 'CircleSoftmax' + "Linear", + "ArcSoftmax", + "CosSoftmax", + "CircleSoftmax" ] class Linear(nn.Module): def __init__(self, num_classes, scale, margin): super().__init__() - self._num_classes = num_classes - self.s = 1 - self.m = 0 + self.num_classes = num_classes + self.s = scale + self.m = margin - def forward(self, logits, *args): + def forward(self, logits, targets): return logits def extra_repr(self): - return 'num_classes={}, scale={}, margin={}'.format(self._num_classes, self.s, self.m) + return f"num_classes={self.num_classes}, scale={self.s}, margin={self.m}" -class ArcSoftmax(nn.Module): - def __init__(self, num_classes, scale, margin): - super().__init__() - self._num_classes = num_classes - self.s = scale - self.m = margin +class CosSoftmax(Linear): + r"""Implement of large margin cosine distance: + """ - self.easy_margin = False + def forward(self, logits, targets): + index = torch.where(targets != -1)[0] + m_hot = torch.zeros(index.size()[0], logits.size()[1], device=logits.device, dtype=logits.dtype) + m_hot.scatter_(1, targets[index, None], self.m) + logits[index] -= m_hot + logits.mul_(self.s) + return logits - self.cos_m = math.cos(self.m) - self.sin_m = math.sin(self.m) - self.threshold = math.cos(math.pi - self.m) - self.mm = math.sin(math.pi - self.m) * self.m + +class ArcSoftmax(Linear): def forward(self, logits, targets): - sine = torch.sqrt(1.0 - torch.pow(logits, 2)) - phi = logits * self.cos_m - sine * self.sin_m # cos(theta + m) - if self.easy_margin: - phi = torch.where(logits > 0, phi, logits) - else: - phi = torch.where(logits > self.threshold, phi, logits - self.mm) - one_hot = torch.zeros(logits.size(), device=logits.device) - one_hot.scatter_(1, targets.view(-1, 1).long(), 1) - output = (one_hot * phi) + ((1.0 - one_hot) * logits) - output *= self.s - return output + index = torch.where(targets != -1)[0] + m_hot = torch.zeros(index.size()[0], logits.size()[1], device=logits.device, dtype=logits.dtype) + m_hot.scatter_(1, targets[index, None], self.m) + logits.acos_() + logits[index] += m_hot + logits.cos_().mul_(self.s) + return logits - def extra_repr(self): - return 'num_classes={}, scale={}, margin={}'.format(self._num_classes, self.s, self.m) - -class CircleSoftmax(nn.Module): - def __init__(self, num_classes, scale, margin): - super().__init__() - self._num_classes = num_classes - self.s = scale - self.m = margin +class CircleSoftmax(Linear): def forward(self, logits, targets): alpha_p = torch.clamp_min(-logits.detach() + 1 + self.m, min=0.) @@ -76,38 +62,19 @@ def forward(self, logits, targets): delta_p = 1 - self.m delta_n = self.m - s_p = self.s * alpha_p * (logits - delta_p) - s_n = self.s * alpha_n * (logits - delta_n) + # When use model parallel, there are some targets not in class centers of local rank + index = torch.where(targets != -1)[0] + m_hot = torch.zeros(index.size()[0], logits.size()[1], device=logits.device, dtype=logits.dtype) + m_hot.scatter_(1, targets[index, None], 1) - targets = F.one_hot(targets, num_classes=self._num_classes) + logits_p = alpha_p * (logits - delta_p) + logits_n = alpha_n * (logits - delta_n) - pred_class_logits = targets * s_p + (1.0 - targets) * s_n + logits[index] = logits_p[index] * m_hot + logits_n[index] * (1 - m_hot) - return pred_class_logits + neg_index = torch.where(targets == -1)[0] + logits[neg_index] = logits_n[neg_index] - def extra_repr(self): - return "num_classes={}, scale={}, margin={}".format(self._num_classes, self.s, self.m) + logits.mul_(self.s) - -class CosSoftmax(nn.Module): - r"""Implement of large margin cosine distance: - Args: - num_classes: size of each output sample - """ - - def __init__(self, num_classes, scale, margin): - super().__init__() - self._num_classes = num_classes - self.s = scale - self.m = margin - - def forward(self, logits, targets): - phi = logits - self.m - targets = F.one_hot(targets, num_classes=self._num_classes) - output = (targets * phi) + ((1.0 - targets) * logits) - output *= self.s - - return output - - def extra_repr(self): - return "num_classes={}, scale={}, margin={}".format(self._num_classes, self.s, self.m) + return logits diff --git a/fastreid/modeling/heads/embedding_head.py b/fastreid/modeling/heads/embedding_head.py index c6a06a79..5fce9006 100644 --- a/fastreid/modeling/heads/embedding_head.py +++ b/fastreid/modeling/heads/embedding_head.py @@ -83,15 +83,7 @@ def __init__( # Linear layer assert hasattr(any_softmax, cls_type), "Expected cls types are {}, " \ "but got {}".format(any_softmax.__all__, cls_type) - self.weight = nn.Parameter(torch.Tensor(num_classes, feat_dim)) - # Initialize weight parameters - if cls_type == "Linear": - nn.init.normal_(self.weight, std=0.001) - elif cls_type == "CircleSoftmax": - nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) - elif cls_type == "ArcSoftmax" or cls_type == "CosSoftmax": - nn.init.xavier_uniform_(self.weigth) - + self.weight = nn.Parameter(torch.normal(0, 0.01, (num_classes, feat_dim))) self.cls_layer = getattr(any_softmax, cls_type)(num_classes, scale, margin) @classmethod diff --git a/projects/FastAttr/fastattr/__init__.py b/projects/FastAttr/fastattr/__init__.py index 93dabc55..fe203238 100644 --- a/projects/FastAttr/fastattr/__init__.py +++ b/projects/FastAttr/fastattr/__init__.py @@ -4,9 +4,9 @@ @contact: sherlockliao01@gmail.com """ -from .attr_baseline import AttrBaseline from .attr_evaluation import AttrEvaluator -from .attr_head import AttrHead from .config import add_attr_config from .data_build import build_attr_train_loader, build_attr_test_loader from .datasets import * +from .modeling import * +from .attr_dataset import AttrDataset diff --git a/projects/FastAttr/fastattr/config.py b/projects/FastAttr/fastattr/config.py index 5b69581c..2ee7f333 100644 --- a/projects/FastAttr/fastattr/config.py +++ b/projects/FastAttr/fastattr/config.py @@ -10,8 +10,7 @@ def add_attr_config(cfg): _C = cfg - _C.MODEL.LOSSES.BCE = CN() - _C.MODEL.LOSSES.BCE.WEIGHT_ENABLED = True + _C.MODEL.LOSSES.BCE = CN({"WEIGHT_ENABLED": True}) _C.MODEL.LOSSES.BCE.SCALE = 1. _C.TEST.THRES = 0.5 diff --git a/projects/FastAttr/fastattr/data_build.py b/projects/FastAttr/fastattr/data_build.py deleted file mode 100644 index d1b9959f..00000000 --- a/projects/FastAttr/fastattr/data_build.py +++ /dev/null @@ -1,74 +0,0 @@ -# encoding: utf-8 -""" -@author: l1aoxingyu -@contact: sherlockliao01@gmail.com -""" - -import os - -import torch -from torch.utils.data import DataLoader - -from fastreid.data import samplers -from fastreid.data.build import fast_batch_collator -from fastreid.data.datasets import DATASET_REGISTRY -from fastreid.data.transforms import build_transforms -from fastreid.utils import comm -from .attr_dataset import AttrDataset - -_root = os.getenv("FASTREID_DATASETS", "datasets") - - -def build_attr_train_loader(cfg): - train_items = list() - attr_dict = None - for d in cfg.DATASETS.NAMES: - dataset = DATASET_REGISTRY.get(d)(root=_root, combineall=cfg.DATASETS.COMBINEALL) - if comm.is_main_process(): - dataset.show_train() - if attr_dict is not None: - assert attr_dict == dataset.attr_dict, f"attr_dict in {d} does not match with previous ones" - else: - attr_dict = dataset.attr_dict - train_items.extend(dataset.train) - - train_transforms = build_transforms(cfg, is_train=True) - train_set = AttrDataset(train_items, train_transforms, attr_dict) - - num_workers = cfg.DATALOADER.NUM_WORKERS - mini_batch_size = cfg.SOLVER.IMS_PER_BATCH // comm.get_world_size() - - data_sampler = samplers.TrainingSampler(len(train_set)) - batch_sampler = torch.utils.data.sampler.BatchSampler(data_sampler, mini_batch_size, True) - - train_loader = torch.utils.data.DataLoader( - train_set, - num_workers=num_workers, - batch_sampler=batch_sampler, - collate_fn=fast_batch_collator, - pin_memory=True, - ) - return train_loader - - -def build_attr_test_loader(cfg, dataset_name): - dataset = DATASET_REGISTRY.get(dataset_name)(root=_root, combineall=cfg.DATASETS.COMBINEALL) - attr_dict = dataset.attr_dict - if comm.is_main_process(): - dataset.show_test() - test_items = dataset.test - - test_transforms = build_transforms(cfg, is_train=False) - test_set = AttrDataset(test_items, test_transforms, attr_dict) - - mini_batch_size = cfg.TEST.IMS_PER_BATCH // comm.get_world_size() - data_sampler = samplers.InferenceSampler(len(test_set)) - batch_sampler = torch.utils.data.BatchSampler(data_sampler, mini_batch_size, False) - test_loader = DataLoader( - test_set, - batch_sampler=batch_sampler, - num_workers=4, # save some memory - collate_fn=fast_batch_collator, - pin_memory=True, - ) - return test_loader diff --git a/projects/FastAttr/fastattr/modeling/__init__.py b/projects/FastAttr/fastattr/modeling/__init__.py new file mode 100644 index 00000000..671b11c6 --- /dev/null +++ b/projects/FastAttr/fastattr/modeling/__init__.py @@ -0,0 +1,9 @@ +# encoding: utf-8 +""" +@author: xingyu liao +@contact: sherlockliao01@gmail.com +""" + +from .attr_baseline import AttrBaseline +from .attr_head import AttrHead +from .bce_loss import cross_entropy_sigmoid_loss diff --git a/projects/FastAttr/fastattr/attr_baseline.py b/projects/FastAttr/fastattr/modeling/attr_baseline.py similarity index 100% rename from projects/FastAttr/fastattr/attr_baseline.py rename to projects/FastAttr/fastattr/modeling/attr_baseline.py diff --git a/projects/FastAttr/fastattr/attr_head.py b/projects/FastAttr/fastattr/modeling/attr_head.py similarity index 100% rename from projects/FastAttr/fastattr/attr_head.py rename to projects/FastAttr/fastattr/modeling/attr_head.py diff --git a/projects/FastAttr/fastattr/bce_loss.py b/projects/FastAttr/fastattr/modeling/bce_loss.py similarity index 100% rename from projects/FastAttr/fastattr/bce_loss.py rename to projects/FastAttr/fastattr/modeling/bce_loss.py diff --git a/projects/FastAttr/train_net.py b/projects/FastAttr/train_net.py index 7eaa0b7a..9a4dcc88 100644 --- a/projects/FastAttr/train_net.py +++ b/projects/FastAttr/train_net.py @@ -3,6 +3,7 @@ @author: xingyu liao @contact: sherlockliao01@gmail.com """ +import logging import sys sys.path.append('.') @@ -11,11 +12,15 @@ from fastreid.engine import DefaultTrainer from fastreid.engine import default_argument_parser, default_setup, launch from fastreid.utils.checkpoint import Checkpointer +from fastreid.data.datasets import DATASET_REGISTRY +from fastreid.data.build import _root, build_reid_train_loader, build_reid_test_loader +from fastreid.data.transforms import build_transforms +from fastreid.utils import comm from fastattr import * -class Trainer(DefaultTrainer): +class AttrTrainer(DefaultTrainer): sample_weights = None @classmethod @@ -28,21 +33,47 @@ def build_model(cls, cfg): """ model = DefaultTrainer.build_model(cfg) if cfg.MODEL.LOSSES.BCE.WEIGHT_ENABLED and \ - Trainer.sample_weights is not None: - setattr(model, "sample_weights", Trainer.sample_weights.to(model.device)) + AttrTrainer.sample_weights is not None: + setattr(model, "sample_weights", AttrTrainer.sample_weights.to(model.device)) else: setattr(model, "sample_weights", None) return model @classmethod def build_train_loader(cls, cfg): - data_loader = build_attr_train_loader(cfg) - Trainer.sample_weights = data_loader.dataset.sample_weights + + logger = logging.getLogger("fastreid.attr_dataset") + train_items = list() + attr_dict = None + for d in cfg.DATASETS.NAMES: + dataset = DATASET_REGISTRY.get(d)(root=_root, combineall=cfg.DATASETS.COMBINEALL) + if comm.is_main_process(): + dataset.show_train() + if attr_dict is not None: + assert attr_dict == dataset.attr_dict, f"attr_dict in {d} does not match with previous ones" + else: + attr_dict = dataset.attr_dict + train_items.extend(dataset.train) + + train_transforms = build_transforms(cfg, is_train=True) + train_set = AttrDataset(train_items, train_transforms, attr_dict) + + data_loader = build_reid_train_loader(cfg, train_set=train_set) + AttrTrainer.sample_weights = data_loader.dataset.sample_weights return data_loader @classmethod def build_test_loader(cls, cfg, dataset_name): - return build_attr_test_loader(cfg, dataset_name) + dataset = DATASET_REGISTRY.get(dataset_name)(root=_root) + attr_dict = dataset.attr_dict + if comm.is_main_process(): + dataset.show_test() + test_items = dataset.test + + test_transforms = build_transforms(cfg, is_train=False) + test_set = AttrDataset(test_items, test_transforms, attr_dict) + data_loader, _ = build_reid_test_loader(cfg, test_set=test_set) + return data_loader @classmethod def build_evaluator(cls, cfg, dataset_name, output_folder=None): @@ -69,14 +100,14 @@ def main(args): if args.eval_only: cfg.defrost() cfg.MODEL.BACKBONE.PRETRAIN = False - model = Trainer.build_model(cfg) + model = AttrTrainer.build_model(cfg) Checkpointer(model).load(cfg.MODEL.WEIGHTS) # load trained model - res = Trainer.test(cfg, model) + res = AttrTrainer.test(cfg, model) return res - trainer = Trainer(cfg) + trainer = AttrTrainer(cfg) trainer.resume_or_load(resume=args.resume) return trainer.train() diff --git a/projects/FastClas/configs/base-clas.yaml b/projects/FastClas/configs/base-clas.yaml index 031ca4cd..bff7e9a3 100644 --- a/projects/FastClas/configs/base-clas.yaml +++ b/projects/FastClas/configs/base-clas.yaml @@ -3,10 +3,10 @@ MODEL: BACKBONE: NAME: build_resnet_backbone - DEPTH: 50x + DEPTH: 18x NORM: BN LAST_STRIDE: 2 - FEAT_DIM: 2048 + FEAT_DIM: 512 PRETRAIN: True HEADS: @@ -25,11 +25,14 @@ MODEL: SCALE: 1. INPUT: - SIZE_TEST: [256, 256] + SIZE_TRAIN: [0,] # no need for resize when training + SIZE_TEST: [256,] CROP: ENABLED: True - SIZE: [224, 224] + SIZE: [224,] + SCALE: [0.08, 1] + RATIO: [0.75, 1.333333333] FLIP: ENABLED: True @@ -46,19 +49,19 @@ SOLVER: OPT: SGD SCHED: CosineAnnealingLR - BASE_LR: 0.003 - MOMENTUM: 0.99 - NESTEROV: True + BASE_LR: 0.001 + MOMENTUM: 0.9 + NESTEROV: False BIAS_LR_FACTOR: 1. WEIGHT_DECAY: 0.0005 WEIGHT_DECAY_BIAS: 0. - IMS_PER_BATCH: 128 + IMS_PER_BATCH: 4 ETA_MIN_LR: 0.00003 WARMUP_FACTOR: 0.1 - WARMUP_ITERS: 2000 + WARMUP_ITERS: 100 CHECKPOINT_PERIOD: 10 @@ -70,4 +73,4 @@ DATASETS: NAMES: ("Hymenoptera",) TESTS: ("Hymenoptera",) -OUTPUT_DIR: projects/FastClas/logs/baseline \ No newline at end of file +OUTPUT_DIR: projects/FastClas/logs/r18_demo \ No newline at end of file diff --git a/projects/FastClas/fastclas/dataset.py b/projects/FastClas/fastclas/dataset.py index 1a46a3d5..7f6ae1a8 100644 --- a/projects/FastClas/fastclas/dataset.py +++ b/projects/FastClas/fastclas/dataset.py @@ -12,7 +12,7 @@ class ClasDataset(Dataset): """Image Person ReID Dataset""" - def __init__(self, img_items, transform=None, relabel=True): + def __init__(self, img_items, transform=None): self.img_items = img_items self.transform = transform diff --git a/projects/FastClas/train_net.py b/projects/FastClas/train_net.py index 9824c4f0..4d779daf 100644 --- a/projects/FastClas/train_net.py +++ b/projects/FastClas/train_net.py @@ -7,8 +7,8 @@ import json import logging -import sys import os +import sys sys.path.append('.') @@ -19,11 +19,14 @@ from fastreid.utils.checkpoint import Checkpointer, PathManager from fastreid.utils import comm from fastreid.engine import DefaultTrainer +from fastreid.data.datasets import DATASET_REGISTRY +from fastreid.data.transforms import build_transforms +from fastreid.data.build import _root from fastclas import * -class Trainer(DefaultTrainer): +class ClasTrainer(DefaultTrainer): @classmethod def build_train_loader(cls, cfg): @@ -35,14 +38,25 @@ def build_train_loader(cls, cfg): """ logger = logging.getLogger("fastreid.clas_dataset") logger.info("Prepare training set") - data_loader = build_reid_train_loader(cfg, Dataset=ClasDataset) + + train_items = list() + for d in cfg.DATASETS.NAMES: + data = DATASET_REGISTRY.get(d)(root=_root) + if comm.is_main_process(): + data.show_train() + train_items.extend(data.train) + + transforms = build_transforms(cfg, is_train=True) + train_set = ClasDataset(train_items, transforms) + + data_loader = build_reid_train_loader(cfg, train_set=train_set) # Save index to class dictionary output_dir = cfg.OUTPUT_DIR if comm.is_main_process() and output_dir: path = os.path.join(output_dir, "idx2class.json") with PathManager.open(path, "w") as f: - json.dump(data_loader.dataset.idx_to_class, f) + json.dump(train_set.idx_to_class, f) return data_loader @@ -54,11 +68,18 @@ def build_test_loader(cls, cfg, dataset_name): It now calls :func:`fastreid.data.build_reid_test_loader`. Overwrite it if you'd like a different data loader. """ - return build_reid_test_loader(cfg, dataset_name, Dataset=ClasDataset) + + data = DATASET_REGISTRY.get(dataset_name)(root=_root) + if comm.is_main_process(): + data.show_test() + transforms = build_transforms(cfg, is_train=False) + test_set = ClasDataset(data.query, transforms) + data_loader, _ = build_reid_test_loader(cfg, test_set=test_set) + return data_loader @classmethod def build_evaluator(cls, cfg, dataset_name, output_dir=None): - data_loader, _ = cls.build_test_loader(cfg, dataset_name) + data_loader = cls.build_test_loader(cfg, dataset_name) return data_loader, ClasEvaluator(cfg, output_dir) @@ -80,14 +101,14 @@ def main(args): if args.eval_only: cfg.defrost() cfg.MODEL.BACKBONE.PRETRAIN = False - model = Trainer.build_model(cfg) + model = ClasTrainer.build_model(cfg) Checkpointer(model).load(cfg.MODEL.WEIGHTS) # load trained model - res = Trainer.test(cfg, model) + res = ClasTrainer.test(cfg, model) return res - trainer = Trainer(cfg) + trainer = ClasTrainer(cfg) trainer.resume_or_load(resume=args.resume) return trainer.train() diff --git a/projects/FastFace/README.md b/projects/FastFace/README.md index 6a8cdfe7..567cbe33 100644 --- a/projects/FastFace/README.md +++ b/projects/FastFace/README.md @@ -20,6 +20,7 @@ We do data wrangling following [InsightFace_Pytorch](https://github.com/TreB1eN/ ## Dependencies - bcolz +- mxnet (optional) if you want to read `.rec` directly ## Experiment Results @@ -29,4 +30,6 @@ We refer to [insightface_pytorch](https://github.com/TreB1eN/InsightFace_Pytorch | :---: | :---: | :---: |:---: | :---: | :---: | :---: | :---: | | [insightface_pytorch](https://github.com/TreB1eN/InsightFace_Pytorch) | 99.52 | 99.62 | 95.04 | 96.22 | 95.57 | 91.07 | 93.86 | | ir50_se | 99.70 | 99.60 | 96.43 | 97.87 | 95.95 | 91.10 | 94.32 | -| ir100_se | 99.65 | 99.69 | 97.10 | 97.98 | 96.00 | 91.53 | 94.62 | \ No newline at end of file +| ir100_se | 99.65 | 99.69 | 97.10 | 97.98 | 96.00 | 91.53 | 94.62 | +| ir50_se_0.1 | | | | | | | | +| ir100_se_0.1 | | | | | | | | diff --git a/projects/FastFace/configs/face_base.yml b/projects/FastFace/configs/face_base.yml index 87402bff..f125cf0c 100644 --- a/projects/FastFace/configs/face_base.yml +++ b/projects/FastFace/configs/face_base.yml @@ -1,18 +1,24 @@ MODEL: - META_ARCHITECTURE: Baseline + META_ARCHITECTURE: FaceBaseline PIXEL_MEAN: [127.5, 127.5, 127.5] PIXEL_STD: [127.5, 127.5, 127.5] HEADS: - NAME: EmbeddingHead + NAME: FaceHead + WITH_BNNECK: True NORM: BN NECK_FEAT: after EMBEDDING_DIM: 512 POOL_LAYER: Flatten - CLS_LAYER: CircleSoftmax - SCALE: 256 - MARGIN: 0.25 + CLS_LAYER: CosSoftmax + SCALE: 64 + MARGIN: 0.4 + NUM_CLASSES: 360232 + + PFC: + ENABLED: False + SAMPLE_RATE: 0.1 LOSSES: NAME: ("CrossEntropyLoss",) @@ -22,33 +28,24 @@ MODEL: SCALE: 1. DATASETS: + REC_PATH: /export/home/DATA/Glint360k/train.rec NAMES: ("MS1MV2",) TESTS: ("CPLFW", "VGG2_FP", "CALFW", "CFP_FF", "CFP_FP", "AgeDB_30", "LFW") INPUT: - SIZE_TRAIN: [112, 112] - SIZE_TEST: [112, 112] - AUGMIX: - ENABLED: False - AUTOAUG: - ENABLED: False - - CJ: - ENABLED: False + SIZE_TRAIN: [0,] # No need of resize + SIZE_TEST: [0,] FLIP: ENABLED: True PROB: 0.5 - PADDING: - ENABLED: False - DATALOADER: SAMPLER_TRAIN: TrainingSampler NUM_WORKERS: 8 SOLVER: - MAX_EPOCH: 16 + MAX_EPOCH: 20 AMP: ENABLED: False @@ -56,8 +53,8 @@ SOLVER: BASE_LR: 0.1 MOMENTUM: 0.9 - SCHED: CosineAnnealingLR - ETA_MIN_LR: 0.0001 + SCHED: MultiStepLR + STEPS: [8, 12, 15, 18] BIAS_LR_FACTOR: 1. WEIGHT_DECAY: 0.0005 @@ -67,10 +64,10 @@ SOLVER: WARMUP_FACTOR: 0.1 WARMUP_ITERS: 5000 - CHECKPOINT_PERIOD: 2 + CHECKPOINT_PERIOD: 1 TEST: - EVAL_PERIOD: 2 + EVAL_PERIOD: 1 IMS_PER_BATCH: 1024 CUDNN_BENCHMARK: True \ No newline at end of file diff --git a/projects/FastFace/configs/r101_ir.yml b/projects/FastFace/configs/r101_ir.yml index a8bc5a31..ac775450 100644 --- a/projects/FastFace/configs/r101_ir.yml +++ b/projects/FastFace/configs/r101_ir.yml @@ -1,7 +1,6 @@ _BASE_: face_base.yml MODEL: - META_ARCHITECTURE: Baseline BACKBONE: NAME: build_resnetIR_backbone @@ -9,4 +8,8 @@ MODEL: FEAT_DIM: 25088 # 512x7x7 WITH_SE: True + HEADS: + PFC: + ENABLED: True + OUTPUT_DIR: projects/FastFace/logs/ir_se101-ms1mv2-circle diff --git a/projects/FastFace/configs/r50_ir.yml b/projects/FastFace/configs/r50_ir.yml index 161d9932..d18bc59d 100644 --- a/projects/FastFace/configs/r50_ir.yml +++ b/projects/FastFace/configs/r50_ir.yml @@ -1,7 +1,6 @@ _BASE_: face_base.yml MODEL: - META_ARCHITECTURE: Baseline BACKBONE: NAME: build_resnetIR_backbone @@ -9,4 +8,8 @@ MODEL: FEAT_DIM: 25088 # 512x7x7 WITH_SE: True -OUTPUT_DIR: projects/FastFace/logs/ir_se50-ms1mv2-circle + HEADS: + PFC: + ENABLED: True + +OUTPUT_DIR: projects/FastFace/logs/ir_se50-glink360k-pfc0.1 diff --git a/projects/FastFace/fastface/__init__.py b/projects/FastFace/fastface/__init__.py index 5895bf0c..10c6a395 100644 --- a/projects/FastFace/fastface/__init__.py +++ b/projects/FastFace/fastface/__init__.py @@ -4,7 +4,6 @@ @contact: sherlockliao01@gmail.com """ -from .datasets import * -from .build import build_face_test_loader -from .resnet_ir import build_resnetIR_backbone -from .face_evaluator import FaceEvaluator +from .modeling import * +from .config import add_face_cfg +from .trainer import FaceTrainer diff --git a/projects/FastFace/fastface/build.py b/projects/FastFace/fastface/build.py deleted file mode 100644 index 97569f22..00000000 --- a/projects/FastFace/fastface/build.py +++ /dev/null @@ -1,46 +0,0 @@ -# encoding: utf-8 -""" -@author: xingyu liao -@contact: sherlockliao01@gmail.com -""" - -import torch -from torch.utils.data import DataLoader - -from fastreid.data import samplers -from fastreid.data.build import fast_batch_collator, _root -from fastreid.data.common import CommDataset -from fastreid.data.datasets import DATASET_REGISTRY -from fastreid.utils import comm - - -class FaceCommDataset(CommDataset): - def __init__(self, img_items, labels): - self.img_items = img_items - self.labels = labels - - def __getitem__(self, index): - img = torch.tensor(self.img_items[index]) * 127.5 + 127.5 - return { - "images": img, - } - - -def build_face_test_loader(cfg, dataset_name, **kwargs): - dataset = DATASET_REGISTRY.get(dataset_name)(root=_root, **kwargs) - if comm.is_main_process(): - dataset.show_test() - - test_set = FaceCommDataset(dataset.carray, dataset.is_same) - - mini_batch_size = cfg.TEST.IMS_PER_BATCH // comm.get_world_size() - data_sampler = samplers.InferenceSampler(len(test_set)) - batch_sampler = torch.utils.data.BatchSampler(data_sampler, mini_batch_size, False) - test_loader = DataLoader( - test_set, - batch_sampler=batch_sampler, - num_workers=4, # save some memory - collate_fn=fast_batch_collator, - pin_memory=True, - ) - return test_loader, test_set.labels diff --git a/projects/FastFace/fastface/config.py b/projects/FastFace/fastface/config.py new file mode 100644 index 00000000..f47fa9f3 --- /dev/null +++ b/projects/FastFace/fastface/config.py @@ -0,0 +1,16 @@ +# encoding: utf-8 +""" +@author: xingyu liao +@contact: sherlockliao01@gmail.com +""" + +from fastreid.config import CfgNode as CN + + +def add_face_cfg(cfg): + _C = cfg + + _C.DATASETS.REC_PATH = "" + + _C.MODEL.HEADS.PFC = CN({"ENABLED": False}) + _C.MODEL.HEADS.PFC.SAMPLE_RATE = 0.1 diff --git a/projects/FastFace/fastface/datasets/ms1mv2.py b/projects/FastFace/fastface/datasets/ms1mv2.py index 75858bf2..b19d4f8b 100644 --- a/projects/FastFace/fastface/datasets/ms1mv2.py +++ b/projects/FastFace/fastface/datasets/ms1mv2.py @@ -21,11 +21,9 @@ def __init__(self, root="datasets", **kwargs): self.dataset_dir = os.path.join(self.root, self.dataset_dir) required_files = [self.dataset_dir] - self.check_before_run(required_files) train = self.process_dirs() - super().__init__(train, [], [], **kwargs) def process_dirs(self): diff --git a/projects/FastFace/fastface/face_data.py b/projects/FastFace/fastface/face_data.py new file mode 100644 index 00000000..cf161ad2 --- /dev/null +++ b/projects/FastFace/fastface/face_data.py @@ -0,0 +1,80 @@ +# encoding: utf-8 +""" +@author: xingyu liao +@contact: sherlockliao01@gmail.com +""" + +from PIL import Image +import io +import logging +import numbers + +import torch +from torch.utils.data import Dataset + +from fastreid.data.common import CommDataset + +logger = logging.getLogger("fastreid.face_data") + +try: + import mxnet as mx +except ImportError: + logger.info("Please install mxnet if you want to use .rec file") + + +class MXFaceDataset(Dataset): + def __init__(self, path_imgrec, transforms): + super().__init__() + self.transforms = transforms + + logger.info(f"loading recordio {path_imgrec}...") + path_imgidx = path_imgrec[0:-4] + ".idx" + self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') + s = self.imgrec.read_idx(0) + header, _ = mx.recordio.unpack(s) + if header.flag > 0: + # logger.debug(f"header0 label: {header.label}") + self.header0 = (int(header.label[0]), int(header.label[1])) + self.imgidx = list(range(1, int(header.label[0]))) + # logger.debug(self.imgidx) + else: + self.imgidx = list(self.imgrec.keys) + logger.info(f"Number of Samples: {len(self.imgidx)}, " + f"Number of Classes: {int(self.header0[1] - self.header0[0])}") + + def __getitem__(self, index): + idx = self.imgidx[index] + s = self.imgrec.read_idx(idx) + header, img = mx.recordio.unpack(s) + label = header.label + if not isinstance(label, numbers.Number): + label = label[0] + label = torch.tensor(label, dtype=torch.long) + + sample = Image.open(io.BytesIO(img)) # RGB + if self.transforms is not None: sample = self.transforms(sample) + return { + "images": sample, + "targets": label, + "camids": 0, + } + + def __len__(self): + # logger.debug(f"mxface dataset length is {len(self.imgidx)}") + return len(self.imgidx) + + @property + def num_classes(self): + return int(self.header0[1] - self.header0[0]) + + +class TestFaceDataset(CommDataset): + def __init__(self, img_items, labels): + self.img_items = img_items + self.labels = labels + + def __getitem__(self, index): + img = torch.tensor(self.img_items[index]) * 127.5 + 127.5 + return { + "images": img, + } diff --git a/projects/FastFace/fastface/modeling/__init__.py b/projects/FastFace/fastface/modeling/__init__.py new file mode 100644 index 00000000..78970146 --- /dev/null +++ b/projects/FastFace/fastface/modeling/__init__.py @@ -0,0 +1,10 @@ +# encoding: utf-8 +""" +@author: xingyu liao +@contact: sherlockliao01@gmail.com +""" + +from .partial_fc import PartialFC +from .face_baseline import FaceBaseline +from .face_head import FaceHead +from .resnet_ir import build_resnetIR_backbone diff --git a/projects/FastFace/fastface/modeling/face_baseline.py b/projects/FastFace/fastface/modeling/face_baseline.py new file mode 100644 index 00000000..9cf6ec78 --- /dev/null +++ b/projects/FastFace/fastface/modeling/face_baseline.py @@ -0,0 +1,24 @@ +# 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 diff --git a/projects/FastFace/fastface/modeling/face_head.py b/projects/FastFace/fastface/modeling/face_head.py new file mode 100644 index 00000000..0168b845 --- /dev/null +++ b/projects/FastFace/fastface/modeling/face_head.py @@ -0,0 +1,39 @@ +# encoding: utf-8 +""" +@author: xingyu liao +@contact: sherlockliao01@gmail.com +""" + +from fastreid.config import configurable +from fastreid.modeling.heads import EmbeddingHead +from fastreid.modeling.heads.build import REID_HEADS_REGISTRY + + +@REID_HEADS_REGISTRY.register() +class FaceHead(EmbeddingHead): + def __init__(self, cfg): + super().__init__(cfg) + self.pfc_enabled = False + if cfg.MODEL.HEADS.PFC.ENABLED: + # Delete pre-defined linear weights for partial fc sample + del self.weight + self.pfc_enabled = True + + def forward(self, features, targets=None): + """ + Partial FC forward, which will sample positive weights and part of negative weights, + then compute logits and get the grad of features. + """ + if not self.pfc_enabled: + return super().forward(features, targets) + else: + pool_feat = self.pool_layer(features) + neck_feat = self.bottleneck(pool_feat) + neck_feat = neck_feat[..., 0, 0] + + if not self.training: + return neck_feat + + return { + "features": neck_feat, + } diff --git a/projects/FastFace/fastface/modeling/partial_fc.py b/projects/FastFace/fastface/modeling/partial_fc.py new file mode 100644 index 00000000..408cace1 --- /dev/null +++ b/projects/FastFace/fastface/modeling/partial_fc.py @@ -0,0 +1,196 @@ +# encoding: utf-8 +# code based on: +# https://github.com/deepinsight/insightface/blob/master/recognition/arcface_torch/partial_fc.py + +import logging +import math + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import nn + +from fastreid.layers import any_softmax +from fastreid.modeling.losses.utils import concat_all_gather +from fastreid.utils import comm + +logger = logging.getLogger('fastreid.partial_fc') + + +class PartialFC(nn.Module): + """ + Author: {Xiang An, Yang Xiao, XuHan Zhu} in DeepGlint, + Partial FC: Training 10 Million Identities on a Single Machine + See the original paper: + https://arxiv.org/abs/2010.05222 + """ + + def __init__( + self, + embedding_size, + num_classes, + sample_rate, + cls_type, + scale, + margin + ): + super().__init__() + + self.embedding_size = embedding_size + self.num_classes = num_classes + self.sample_rate = sample_rate + + self.world_size = comm.get_world_size() + self.rank = comm.get_rank() + self.local_rank = comm.get_local_rank() + self.device = torch.device(f'cuda:{self.local_rank}') + + self.num_local: int = self.num_classes // self.world_size + int(self.rank < self.num_classes % self.world_size) + self.class_start: int = self.num_classes // self.world_size * self.rank + \ + min(self.rank, self.num_classes % self.world_size) + self.num_sample: int = int(self.sample_rate * self.num_local) + + self.cls_layer = getattr(any_softmax, cls_type)(num_classes, scale, margin) + + """ TODO: consider resume training + if resume: + try: + self.weight: torch.Tensor = torch.load(self.weight_name) + logging.info("softmax weight resume successfully!") + except (FileNotFoundError, KeyError, IndexError): + self.weight = torch.normal(0, 0.01, (self.num_local, self.embedding_size), device=self.device) + logging.info("softmax weight resume fail!") + + try: + self.weight_mom: torch.Tensor = torch.load(self.weight_mom_name) + logging.info("softmax weight mom resume successfully!") + except (FileNotFoundError, KeyError, IndexError): + self.weight_mom: torch.Tensor = torch.zeros_like(self.weight) + logging.info("softmax weight mom resume fail!") + else: + """ + self.weight = torch.normal(0, 0.01, (self.num_local, self.embedding_size), device=self.device) + self.weight_mom: torch.Tensor = torch.zeros_like(self.weight) + logger.info("softmax weight init successfully!") + logger.info("softmax weight mom init successfully!") + self.stream: torch.cuda.Stream = torch.cuda.Stream(self.local_rank) + + self.index = None + if int(self.sample_rate) == 1: + self.update = lambda: 0 + self.sub_weight = nn.Parameter(self.weight) + self.sub_weight_mom = self.weight_mom + else: + self.sub_weight = nn.Parameter(torch.empty((0, 0), device=self.device)) + + def forward(self, total_features): + torch.cuda.current_stream().wait_stream(self.stream) + if self.cls_layer.__class__.__name__ == 'Linear': + logits = F.linear(total_features, self.sub_weight) + else: + logits = F.linear(F.normalize(total_features), F.normalize(self.sub_weight)) + return logits + + def forward_backward(self, features, targets, optimizer): + """ + Partial FC forward, which will sample positive weights and part of negative weights, + then compute logits and get the grad of features. + """ + total_targets = self.prepare(targets, optimizer) + + if self.world_size > 1: + total_features = concat_all_gather(features) + else: + total_features = features.detach() + + total_features.requires_grad_(True) + + logits = self.forward(total_features) + logits = self.cls_layer(logits, total_targets) + + # from ipdb import set_trace; set_trace() + with torch.no_grad(): + max_fc = torch.max(logits, dim=1, keepdim=True)[0] + if self.world_size > 1: + dist.all_reduce(max_fc, dist.ReduceOp.MAX) + + # calculate exp(logits) and all-reduce + logits_exp = torch.exp(logits - max_fc) + logits_sum_exp = logits_exp.sum(dim=1, keepdim=True) + + if self.world_size > 1: + dist.all_reduce(logits_sum_exp, dist.ReduceOp.SUM) + + # calculate prob + logits_exp.div_(logits_sum_exp) + + # get one-hot + grad = logits_exp + index = torch.where(total_targets != -1)[0] + one_hot = torch.zeros(size=[index.size()[0], grad.size()[1]], device=grad.device) + one_hot.scatter_(1, total_targets[index, None], 1) + + # calculate loss + loss = torch.zeros(grad.size()[0], 1, device=grad.device) + loss[index] = grad[index].gather(1, total_targets[index, None]) + if self.world_size > 1: + dist.all_reduce(loss, dist.ReduceOp.SUM) + loss_v = loss.clamp_min_(1e-30).log_().mean() * (-1) + + # calculate grad + grad[index] -= one_hot + grad.div_(logits.size(0)) + + logits.backward(grad) + if total_features.grad is not None: + total_features.grad.detach_() + x_grad: torch.Tensor = torch.zeros_like(features) + # feature gradient all-reduce + if self.world_size > 1: + dist.reduce_scatter(x_grad, list(total_features.grad.chunk(self.world_size, dim=0))) + else: + x_grad = total_features.grad + x_grad = x_grad * self.world_size + # backward backbone + return x_grad, loss_v + + @torch.no_grad() + def sample(self, total_targets): + """ + Get sub_weights according to total targets gathered from all GPUs, due to each weights in different + GPU contains different class centers. + """ + index_positive = (self.class_start <= total_targets) & (total_targets < self.class_start + self.num_local) + total_targets[~index_positive] = -1 + total_targets[index_positive] -= self.class_start + if int(self.sample_rate) != 1: + positive = torch.unique(total_targets[index_positive], sorted=True) + if self.num_sample - positive.size(0) >= 0: + perm = torch.rand(size=[self.num_local], device=self.weight.device) + perm[positive] = 2.0 + index = torch.topk(perm, k=self.num_sample)[1] + index = index.sort()[0] + else: + index = positive + self.index = index + total_targets[index_positive] = torch.searchsorted(index, total_targets[index_positive]) + self.sub_weight = nn.Parameter(self.weight[index]) + self.sub_weight_mom = self.weight_mom[index] + + @torch.no_grad() + def update(self): + self.weight_mom[self.index] = self.sub_weight_mom + self.weight[self.index] = self.sub_weight + + def prepare(self, targets, optimizer): + with torch.cuda.stream(self.stream): + if self.world_size > 1: + total_targets = concat_all_gather(targets) + else: + total_targets = targets + # update sub_weight + self.sample(total_targets) + optimizer.state.pop(optimizer.param_groups[-1]['params'][0], None) + optimizer.param_groups[-1]['params'][0] = self.sub_weight + optimizer.state[self.sub_weight]["momentum_buffer"] = self.sub_weight_mom + return total_targets diff --git a/projects/FastFace/fastface/resnet_ir.py b/projects/FastFace/fastface/modeling/resnet_ir.py similarity index 100% rename from projects/FastFace/fastface/resnet_ir.py rename to projects/FastFace/fastface/modeling/resnet_ir.py diff --git a/projects/FastFace/fastface/trainer.py b/projects/FastFace/fastface/trainer.py new file mode 100644 index 00000000..b38a5046 --- /dev/null +++ b/projects/FastFace/fastface/trainer.py @@ -0,0 +1,173 @@ +# encoding: utf-8 +""" +@author: xingyu liao +@contact: sherlockliao01@gmail.com +""" +import logging +import os +import time + +from torch.nn.parallel import DistributedDataParallel + +from fastreid.engine import hooks +from .face_data import TestFaceDataset +from fastreid.data.datasets import DATASET_REGISTRY +from fastreid.data.build import _root, build_reid_test_loader, build_reid_train_loader +from fastreid.data.transforms import build_transforms +from fastreid.engine.defaults import DefaultTrainer, TrainerBase +from fastreid.engine.train_loop import SimpleTrainer +from fastreid.utils import comm +from fastreid.utils.checkpoint import Checkpointer +from fastreid.utils.logger import setup_logger +from .face_data import MXFaceDataset +from .face_evaluator import FaceEvaluator +from .modeling import PartialFC + + +class FaceTrainer(DefaultTrainer): + def __init__(self, cfg): + TrainerBase.__init__(self) + + logger = logging.getLogger('fastreid.partial-fc.trainer') + if not logger.isEnabledFor(logging.INFO): # setup_logger is not called for fastreid + setup_logger() + + # Assume these objects must be constructed in this order. + data_loader = self.build_train_loader(cfg) + cfg = self.auto_scale_hyperparams(cfg, data_loader.dataset.num_classes) + model = self.build_model(cfg) + optimizer = self.build_optimizer(cfg, model) + + if cfg.MODEL.HEADS.PFC.ENABLED: + # fmt: off + feat_dim = cfg.MODEL.BACKBONE.FEAT_DIM + embedding_dim = cfg.MODEL.HEADS.EMBEDDING_DIM + num_classes = cfg.MODEL.HEADS.NUM_CLASSES + sample_rate = cfg.MODEL.HEADS.PFC.SAMPLE_RATE + cls_type = cfg.MODEL.HEADS.CLS_LAYER + scale = cfg.MODEL.HEADS.SCALE + margin = cfg.MODEL.HEADS.MARGIN + # fmt: on + # Partial-FC module + embedding_size = embedding_dim if embedding_dim > 0 else feat_dim + self.pfc_module = PartialFC(embedding_size, num_classes, sample_rate, cls_type, scale, margin) + self.pfc_optimizer = self.build_optimizer(cfg, self.pfc_module) + + # For training, wrap with DDP. But don't need this for inference. + if comm.get_world_size() > 1: + # ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True` + # for part of the parameters is not updated. + model = DistributedDataParallel( + model, device_ids=[comm.get_local_rank()], broadcast_buffers=False, + find_unused_parameters=True + ) + + self._trainer = PFCTrainer(model, data_loader, optimizer, self.pfc_module, self.pfc_optimizer) \ + if cfg.MODEL.HEADS.PFC.ENABLED else SimpleTrainer(model, data_loader, optimizer) + + self.iters_per_epoch = len(data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH + self.scheduler = self.build_lr_scheduler(cfg, optimizer, self.iters_per_epoch) + if cfg.MODEL.HEADS.PFC.ENABLED: + self.pfc_scheduler = self.build_lr_scheduler(cfg, self.pfc_optimizer, self.iters_per_epoch) + + self.checkpointer = Checkpointer( + # Assume you want to save checkpoints together with logs/statistics + model, + cfg.OUTPUT_DIR, + save_to_disk=comm.is_main_process(), + optimizer=optimizer, + **self.scheduler, + ) + + if cfg.MODEL.HEADS.PFC.ENABLED: + self.pfc_checkpointer = Checkpointer( + self.pfc_module, + cfg.OUTPUT_DIR, + optimizer=self.pfc_optimizer, + **self.pfc_scheduler, + ) + + self.start_epoch = 0 + self.max_epoch = cfg.SOLVER.MAX_EPOCH + self.max_iter = self.max_epoch * self.iters_per_epoch + self.warmup_iters = cfg.SOLVER.WARMUP_ITERS + self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS + self.cfg = cfg + + self.register_hooks(self.build_hooks()) + + def build_hooks(self): + ret = super().build_hooks() + + if self.cfg.MODEL.HEADS.PFC.ENABLED: + # partial fc scheduler hook + ret.append( + hooks.LRScheduler(self.pfc_optimizer, self.pfc_scheduler) + ) + return ret + + @classmethod + def build_train_loader(cls, cfg): + path_imgrec = cfg.DATASETS.REC_PATH + if path_imgrec is not "": + transforms = build_transforms(cfg, is_train=True) + train_set = MXFaceDataset(path_imgrec, transforms) + return build_reid_train_loader(cfg, train_set=train_set) + else: + return DefaultTrainer.build_train_loader(cfg) + + @classmethod + def build_test_loader(cls, cfg, dataset_name): + dataset = DATASET_REGISTRY.get(dataset_name)(root=_root) + test_set = TestFaceDataset(dataset.carray, dataset.is_same) + data_loader, _ = build_reid_test_loader(cfg, test_set=test_set) + return data_loader, test_set.labels + + @classmethod + def build_evaluator(cls, cfg, dataset_name, output_dir=None): + if output_dir is None: + output_dir = os.path.join(cfg.OUTPUT_DIR, "visualization") + data_loader, labels = cls.build_test_loader(cfg, dataset_name) + return data_loader, FaceEvaluator(cfg, labels, dataset_name, output_dir) + + +class PFCTrainer(SimpleTrainer): + """ + Author: {Xiang An, Yang Xiao, XuHan Zhu} in DeepGlint, + Partial FC: Training 10 Million Identities on a Single Machine + See the original paper: + https://arxiv.org/abs/2010.05222 + code based on: + https://github.com/deepinsight/insightface/blob/master/recognition/arcface_torch/partial_fc.py + """ + + def __init__(self, model, data_loader, optimizer, pfc_module, pfc_optimizer): + super().__init__(model, data_loader, optimizer) + + self.pfc_module = pfc_module + self.pfc_optimizer = pfc_optimizer + + def run_step(self): + assert self.model.training, "[PFCTrainer] model was changed to eval mode!" + start = time.perf_counter() + + data = next(self._data_loader_iter) + data_time = time.perf_counter() - start + + features, targets = self.model(data) + + self.optimizer.zero_grad() + self.pfc_optimizer.zero_grad() + + # Partial-fc backward + f_grad, loss_v = self.pfc_module.forward_backward(features, targets, self.pfc_optimizer) + + features.backward(f_grad) + + loss_dict = {"loss_cls": loss_v} + self._write_metrics(loss_dict, data_time) + + self.optimizer.step() + self.pfc_optimizer.step() + + self.pfc_module.update() diff --git a/projects/FastFace/train_net.py b/projects/FastFace/train_net.py index 1cce4075..2b474011 100644 --- a/projects/FastFace/train_net.py +++ b/projects/FastFace/train_net.py @@ -5,36 +5,16 @@ @contact: sherlockliao01@gmail.com """ -import os import sys sys.path.append('.') from fastreid.config import get_cfg -from fastreid.engine import DefaultTrainer, default_argument_parser, default_setup, launch +from fastreid.engine import default_argument_parser, default_setup, launch from fastreid.utils.checkpoint import Checkpointer from fastface import * - - -class Trainer(DefaultTrainer): - - @classmethod - def build_test_loader(cls, cfg, dataset_name): - """ - Returns: - iterable - It now calls :func:`fastreid.data.build_detection_test_loader`. - Overwrite it if you'd like a different data loader. - """ - return build_face_test_loader(cfg, dataset_name) - - @classmethod - def build_evaluator(cls, cfg, dataset_name, output_dir=None): - if output_dir is None: - output_dir = os.path.join(cfg.OUTPUT_DIR, "visualization") - data_loader, labels = cls.build_test_loader(cfg, dataset_name) - return data_loader, FaceEvaluator(cfg, labels, dataset_name, output_dir) +from fastface.datasets import * def setup(args): @@ -42,6 +22,7 @@ def setup(args): Create configs and perform basic setups. """ cfg = get_cfg() + add_face_cfg(cfg) cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() @@ -55,14 +36,14 @@ def main(args): if args.eval_only: cfg.defrost() cfg.MODEL.BACKBONE.PRETRAIN = False - model = Trainer.build_model(cfg) + model = FaceTrainer.build_model(cfg) Checkpointer(model).load(cfg.MODEL.WEIGHTS) # load trained model - res = Trainer.test(cfg, model) + res = FaceTrainer.test(cfg, model) return res - trainer = Trainer(cfg) + trainer = FaceTrainer(cfg) trainer.resume_or_load(resume=args.resume) return trainer.train() diff --git a/projects/PartialReID/configs/partial_market.yml b/projects/PartialReID/configs/partial_market.yml index d43a4e5d..a5613114 100644 --- a/projects/PartialReID/configs/partial_market.yml +++ b/projects/PartialReID/configs/partial_market.yml @@ -1,70 +1,74 @@ MODEL: - META_ARCHITECTURE: 'PartialBaseline' + META_ARCHITECTURE: PartialBaseline BACKBONE: - NAME: "build_resnet_backbone" - DEPTH: "50x" - NORM: "BN" + NAME: build_resnet_backbone + NORM: BN + DEPTH: 50x LAST_STRIDE: 1 + FEAT_DIM: 2048 WITH_IBN: True - PRETRAIN_PATH: "/export/home/lxy/.cache/torch/checkpoints/resnet50_ibn_a-d9d0bb7b.pth" + PRETRAIN: True HEADS: - NAME: "DSRHead" - NORM: "BN" - POOL_LAYER: "avgpool" - NECK_FEAT: "before" - CLS_LAYER: "linear" + NAME: DSRHead + POOL_LAYER: FastGlobalAvgPool + WITH_BNNECK: True + CLS_LAYER: Linear LOSSES: - NAME: ("CrossEntropyLoss", "TripletLoss") + NAME: ("CrossEntropyLoss", "TripletLoss",) + CE: - EPSILON: 0.1 + EPSILON: 0.12 SCALE: 1. + TRI: MARGIN: 0.3 - HARD_MINING: False - SCALE: 1. + SCALE: 1.0 + HARD_MINING: True DATASETS: NAMES: ("Market1501",) - TESTS: ("PartialREID", "PartialiLIDS","OccludedREID",) + TESTS: ("PartialREID", "PartialiLIDS", "OccludedREID",) INPUT: SIZE_TRAIN: [384, 128] SIZE_TEST: [384, 128] - REA: - ENABLED: False - DO_PAD: False + + FLIP: + ENABLED: True DATALOADER: - PK_SAMPLER: True - NAIVE_WAY: False + SAMPLER_TRAIN: NaiveIdentitySampler NUM_INSTANCE: 4 NUM_WORKERS: 8 + SOLVER: - OPT: "Adam" - MAX_ITER: 30 - BASE_LR: 0.00035 - BIAS_LR_FACTOR: 2. + AMP: + ENABLED: False + OPT: Adam + MAX_EPOCH: 60 + BASE_LR: 0.0007 + BIAS_LR_FACTOR: 1. WEIGHT_DECAY: 0.0005 - WEIGHT_DECAY_BIAS: 0.0 - IMS_PER_BATCH: 64 + WEIGHT_DECAY_BIAS: 0.0005 + IMS_PER_BATCH: 256 - SCHED: "WarmupMultiStepLR" - STEPS: [15, 25] - GAMMA: 0.1 + SCHED: CosineAnnealingLR + DELAY_EPOCHS: 20 + ETA_MIN_LR: 0.0000007 - WARMUP_FACTOR: 0.01 - WARMUP_ITERS: 1000 + WARMUP_FACTOR: 0.1 + WARMUP_ITERS: 500 - CHECKPOINT_PERIOD: 10 + CHECKPOINT_PERIOD: 20 TEST: - EVAL_PERIOD: 5 + EVAL_PERIOD: 10 IMS_PER_BATCH: 128 CUDNN_BENCHMARK: True -OUTPUT_DIR: "projects/PartialReID/logs/test_partial" +OUTPUT_DIR: "projects/PartialReID/logs/test_partial" \ No newline at end of file diff --git a/projects/PartialReID/partialreid/config.py b/projects/PartialReID/partialreid/config.py index 882be130..86836582 100644 --- a/projects/PartialReID/partialreid/config.py +++ b/projects/PartialReID/partialreid/config.py @@ -10,6 +10,4 @@ def add_partialreid_config(cfg): _C = cfg - _C.TEST.DSR = CN() - _C.TEST.DSR.ENABLED = True - + _C.TEST.DSR = CN({"ENABLED": True}) diff --git a/projects/PartialReID/partialreid/dsr_distance.py b/projects/PartialReID/partialreid/dsr_distance.py index 4b05d10c..1caedf6a 100644 --- a/projects/PartialReID/partialreid/dsr_distance.py +++ b/projects/PartialReID/partialreid/dsr_distance.py @@ -2,8 +2,8 @@ Notice the input/output shape of methods, so that you can better understand the meaning of these methods.""" -import torch import numpy as np +import torch def normalize(nparray, order=2, axis=0): @@ -46,7 +46,7 @@ def compute_dsr_dist(array1, array2, distmat, scores): Proj_M = torch.FloatTensor(M[index[i, j]]) Proj_M = Proj_M.cuda() a = torch.matmul(g, torch.matmul(Proj_M, q)) - q - dist[i, index[i, j]] = ((torch.pow(a, 2).sum(0).sqrt()) * scores[i]).sum() + dist[i, index[i, j]] = ((torch.pow(a, 2).sum(0).sqrt()) * scores[i].cuda()).sum() dist = dist.cpu() dist = dist.numpy() diff --git a/projects/PartialReID/partialreid/dsr_evaluation.py b/projects/PartialReID/partialreid/dsr_evaluation.py index 9a6c2663..74ec373e 100644 --- a/projects/PartialReID/partialreid/dsr_evaluation.py +++ b/projects/PartialReID/partialreid/dsr_evaluation.py @@ -10,11 +10,9 @@ import numpy as np import torch import torch.nn.functional as F -from sklearn import metrics from fastreid.evaluation.evaluator import DatasetEvaluator from fastreid.evaluation.rank import evaluate_rank -from fastreid.evaluation.roc import evaluate_roc from fastreid.utils import comm from .dsr_distance import compute_dsr_dist @@ -46,7 +44,7 @@ def process(self, inputs, outputs): self.features.append(F.normalize(outputs[0]).cpu()) outputs1 = F.normalize(outputs[1].data).cpu() self.spatial_features.append(outputs1) - self.scores.append(outputs[2]) + self.scores.append(outputs[2].cpu()) def evaluate(self): if comm.get_world_size() > 1: @@ -101,28 +99,28 @@ def evaluate(self): gallery_features = gallery_features.numpy() if self.cfg.TEST.DSR.ENABLED: logger.info("Testing with DSR setting") - dist = compute_dsr_dist(spatial_features[:self._num_query], spatial_features[self._num_query:], dist, - scores[:self._num_query]) - cmc, all_AP, all_INP = evaluate_rank(dist, query_features, gallery_features, query_pids, gallery_pids, - query_camids, gallery_camids, use_distmat=True) + dsr_dist = compute_dsr_dist(spatial_features[:self._num_query], spatial_features[self._num_query:], dist, + scores[:self._num_query]) + + max_value = 0 + k = 0 + for i in range(0, 101): + lamb = 0.01 * i + dist1 = (1 - lamb) * dist + lamb * dsr_dist + cmc, all_AP, all_INP = evaluate_rank(dist1, query_pids, gallery_pids, query_camids, gallery_camids) + if (cmc[0] > max_value): + k = lamb + max_value = cmc[0] + dist1 = (1 - k) * dist + k * dsr_dist + cmc, all_AP, all_INP = evaluate_rank(dist1, query_pids, gallery_pids, query_camids, gallery_camids) else: - cmc, all_AP, all_INP = evaluate_rank(dist, query_features, gallery_features, query_pids, gallery_pids, - query_camids, gallery_camids, use_distmat=False) + cmc, all_AP, all_INP = evaluate_rank(dist, query_pids, gallery_pids, query_camids, gallery_camids) + mAP = np.mean(all_AP) mINP = np.mean(all_INP) - for r in [1, 5, 10]: - self._results['Rank-{}'.format(r)] = cmc[r - 1] - self._results['mAP'] = mAP - self._results['mINP'] = mINP - - if self.cfg.TEST.ROC_ENABLED: - scores, labels = evaluate_roc(dist, query_features, gallery_features, - query_pids, gallery_pids, query_camids, gallery_camids) - fprs, tprs, thres = metrics.roc_curve(labels, scores) - - for fpr in [1e-4, 1e-3, 1e-2]: - ind = np.argmin(np.abs(fprs - fpr)) - self._results["TPR@FPR={:.0e}".format(fpr)] = tprs[ind] + self._results['Rank-{}'.format(r)] = cmc[r - 1] * 100 + self._results['mAP'] = mAP * 100 + self._results['mINP'] = mINP * 100 return copy.deepcopy(self._results) diff --git a/projects/PartialReID/partialreid/dsr_head.py b/projects/PartialReID/partialreid/dsr_head.py index 70091802..57b03ced 100644 --- a/projects/PartialReID/partialreid/dsr_head.py +++ b/projects/PartialReID/partialreid/dsr_head.py @@ -9,8 +9,9 @@ from torch import nn from fastreid.layers import * +from fastreid.modeling.heads import EmbeddingHead from fastreid.modeling.heads.build import REID_HEADS_REGISTRY -from fastreid.utils.weight_init import weights_init_classifier, weights_init_kaiming +from fastreid.utils.weight_init import weights_init_kaiming class OcclusionUnit(nn.Module): @@ -20,7 +21,7 @@ def __init__(self, in_planes=2048): self.MaxPool2 = nn.MaxPool2d(kernel_size=4, stride=2, padding=0) self.MaxPool3 = nn.MaxPool2d(kernel_size=6, stride=2, padding=0) self.MaxPool4 = nn.MaxPool2d(kernel_size=8, stride=2, padding=0) - self.mask_layer = nn.Linear(in_planes, 1, bias=False) + self.mask_layer = nn.Linear(in_planes, 1, bias=True) def forward(self, x): SpaFeat1 = self.MaxPool1(x) # shape: [n, c, h, w] @@ -36,10 +37,13 @@ def forward(self, x): SpatialFeatAll = SpatialFeatAll.transpose(1, 2) # shape: [n, c, m] y = self.mask_layer(SpatialFeatAll) mask_weight = torch.sigmoid(y[:, :, 0]) - + # mask_score = torch.sigmoid(mask_weight[:, :48]) feat_dim = SpaFeat1.size(2) * SpaFeat1.size(3) mask_score = F.normalize(mask_weight[:, :feat_dim], p=1, dim=1) + # mask_score_norm = mask_score + # mask_weight_norm = torch.sigmoid(mask_weight) mask_weight_norm = F.normalize(mask_weight, p=1, dim=1) + mask_score = mask_score.unsqueeze(1) SpaFeat1 = SpaFeat1.transpose(1, 2) @@ -50,61 +54,39 @@ def forward(self, x): return global_feats, mask_weight, mask_weight_norm +class Flatten(nn.Module): + def forward(self, input): + return input.view(input.size(0), -1) + + @REID_HEADS_REGISTRY.register() -class DSRHead(nn.Module): +class DSRHead(EmbeddingHead): def __init__(self, cfg): - super().__init__() - - # fmt: off - feat_dim = cfg.MODEL.BACKBONE.FEAT_DIM - num_classes = cfg.MODEL.HEADS.NUM_CLASSES - neck_feat = cfg.MODEL.HEADS.NECK_FEAT - pool_type = cfg.MODEL.HEADS.POOL_LAYER - cls_type = cfg.MODEL.HEADS.CLS_LAYER - norm_type = cfg.MODEL.HEADS.NORM - - if pool_type == 'fastavgpool': self.pool_layer = FastGlobalAvgPool2d() - elif pool_type == 'avgpool': self.pool_layer = nn.AdaptiveAvgPool2d(1) - elif pool_type == 'maxpool': self.pool_layer = nn.AdaptiveMaxPool2d(1) - elif pool_type == 'gempoolP': self.pool_layer = GeneralizedMeanPoolingP() - elif pool_type == 'gempool': self.pool_layer = GeneralizedMeanPooling() - elif pool_type == "avgmaxpool": self.pool_layer = AdaptiveAvgMaxPool2d() - elif pool_type == 'clipavgpool': self.pool_layer = ClipGlobalAvgPool2d() - elif pool_type == "identity": self.pool_layer = nn.Identity() - elif pool_type == "flatten": self.pool_layer = Flatten() - else: raise KeyError(f"{pool_type} is not supported!") - # fmt: on - - self.neck_feat = neck_feat + super().__init__(cfg) + feat_dim = cfg.MODEL.BACKBONE.FEAT_DIM + with_bnneck = cfg.MODEL.HEADS.WITH_BNNECK + norm_type = cfg.MODEL.HEADS.NORM + num_classes = cfg.MODEL.HEADS.NUM_CLASSES + embedding_dim = cfg.MODEL.HEADS.EMBEDDING_DIM self.occ_unit = OcclusionUnit(in_planes=feat_dim) self.MaxPool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0) self.MaxPool2 = nn.MaxPool2d(kernel_size=4, stride=2, padding=0) self.MaxPool3 = nn.MaxPool2d(kernel_size=6, stride=2, padding=0) self.MaxPool4 = nn.MaxPool2d(kernel_size=8, stride=2, padding=0) - self.bnneck = get_norm(norm_type, feat_dim, bias_freeze=True) - self.bnneck.apply(weights_init_kaiming) + occ_neck = [] + if embedding_dim > 0: + occ_neck.append(nn.Conv2d(feat_dim, embedding_dim, 1, 1, bias=False)) + feat_dim = embedding_dim - self.bnneck_occ = get_norm(norm_type, feat_dim, bias_freeze=True) - self.bnneck_occ.apply(weights_init_kaiming) + if with_bnneck: + occ_neck.append(get_norm(norm_type, feat_dim, bias_freeze=True)) - # identity classification layer - if cls_type == 'linear': - self.classifier = nn.Linear(feat_dim, num_classes, bias=False) - self.classifier_occ = nn.Linear(feat_dim, num_classes, bias=False) - elif cls_type == 'arcSoftmax': - self.classifier = ArcSoftmax(cfg, feat_dim, num_classes) - self.classifier_occ = ArcSoftmax(cfg, feat_dim, num_classes) - elif cls_type == 'circleSoftmax': - self.classifier = CircleSoftmax(cfg, feat_dim, num_classes) - self.classifier_occ = CircleSoftmax(cfg, feat_dim, num_classes) - else: - raise KeyError(f"{cls_type} is invalid, please choose from " - f"'linear', 'arcSoftmax' and 'circleSoftmax'.") + self.bnneck_occ = nn.Sequential(*occ_neck) + self.bnneck_occ.apply(weights_init_kaiming) - self.classifier.apply(weights_init_classifier) - self.classifier_occ.apply(weights_init_classifier) + self.weight_occ = nn.Parameter(torch.normal(0, 0.01, (num_classes, feat_dim))) def forward(self, features, targets=None): """ @@ -122,6 +104,7 @@ def forward(self, features, targets=None): SpatialFeatAll = torch.cat((Feat1, Feat2, Feat3, Feat4), dim=2) foreground_feat, mask_weight, mask_weight_norm = self.occ_unit(features) + # print(time.time() - st) bn_foreground_feat = self.bnneck_occ(foreground_feat) bn_foreground_feat = bn_foreground_feat[..., 0, 0] @@ -131,23 +114,24 @@ def forward(self, features, targets=None): # Training global_feat = self.pool_layer(features) - bn_feat = self.bnneck(global_feat) + bn_feat = self.bottleneck(global_feat) bn_feat = bn_feat[..., 0, 0] - if self.classifier.__class__.__name__ == 'Linear': - cls_outputs = self.classifier(bn_feat) - fore_cls_outputs = self.classifier_occ(bn_foreground_feat) - pred_class_logits = F.linear(bn_feat, self.classifier.weight) + if self.cls_layer.__class__.__name__ == 'Linear': + pred_class_logits = F.linear(bn_feat, self.weight) + fore_pred_class_logits = F.linear(bn_foreground_feat, self.weight_occ) else: - cls_outputs = self.classifier(bn_feat, targets) - fore_cls_outputs = self.classifier_occ(bn_foreground_feat, targets) - pred_class_logits = self.classifier.s * F.linear(F.normalize(bn_feat), - F.normalize(self.classifier.weight)) + pred_class_logits = F.linear(F.normalize(bn_feat), F.normalize(self.weight)) + fore_pred_class_logits = F.linear(F.normalize(bn_foreground_feat), F.normalize(self.weight_occ)) + + cls_outputs = self.cls_layer(pred_class_logits, targets) + fore_cls_outputs = self.cls_layer(fore_pred_class_logits, targets) + # pdb.set_trace() return { "cls_outputs": cls_outputs, "fore_cls_outputs": fore_cls_outputs, - "pred_class_logits": pred_class_logits, - "global_features": global_feat[..., 0, 0], + "pred_class_logits": pred_class_logits * self.cls_layer.s, + "features": global_feat[..., 0, 0], "foreground_features": foreground_feat[..., 0, 0], } diff --git a/projects/PartialReID/partialreid/partialbaseline.py b/projects/PartialReID/partialreid/partialbaseline.py index 088df16e..c4ab4aa1 100644 --- a/projects/PartialReID/partialreid/partialbaseline.py +++ b/projects/PartialReID/partialreid/partialbaseline.py @@ -12,58 +12,35 @@ @META_ARCH_REGISTRY.register() class PartialBaseline(Baseline): - def losses(self, outs): + def losses(self, outputs, gt_labels): r""" Compute loss from modeling's outputs, the loss function input arguments must be the same as the outputs of the model forwarding. """ - # fmt: off - outputs = outs["outputs"] - gt_labels = outs["targets"] - # model predictions - pred_class_logits = outputs['pred_class_logits'].detach() - cls_outputs = outputs["cls_outputs"] - fore_cls_outputs = outputs["fore_cls_outputs"] - global_feat = outputs["global_features"] - fore_feat = outputs["foreground_features"] - # fmt: on + loss_dict = super().losses(outputs, gt_labels) - # Log prediction accuracy - log_accuracy(pred_class_logits, gt_labels) + fore_cls_outputs = outputs["fore_cls_outputs"] + fore_feat = outputs["foreground_features"] - loss_dict = {} - loss_names = self._cfg.MODEL.LOSSES.NAME + loss_names = self.loss_kwargs['loss_names'] - if "CrossEntropyLoss" in loss_names: - loss_dict['loss_avg_branch_cls'] = cross_entropy_loss( - cls_outputs, - gt_labels, - self._cfg.MODEL.LOSSES.CE.EPSILON, - self._cfg.MODEL.LOSSES.CE.ALPHA, - ) * self._cfg.MODEL.LOSSES.CE.SCALE - - loss_dict['loss_fore_branch_cls'] = cross_entropy_loss( + if 'CrossEntropyLoss' in loss_names: + ce_kwargs = self.loss_kwargs.get('ce') + loss_dict['loss_fore_cls'] = cross_entropy_loss( fore_cls_outputs, gt_labels, - self._cfg.MODEL.LOSSES.CE.EPSILON, - self._cfg.MODEL.LOSSES.CE.ALPHA, - ) * self._cfg.MODEL.LOSSES.CE.SCALE - - if "TripletLoss" in loss_names: - loss_dict['loss_avg_branch_triplet'] = triplet_loss( - global_feat, - gt_labels, - self._cfg.MODEL.LOSSES.TRI.MARGIN, - self._cfg.MODEL.LOSSES.TRI.NORM_FEAT, - self._cfg.MODEL.LOSSES.TRI.HARD_MINING, - ) * self._cfg.MODEL.LOSSES.TRI.SCALE + ce_kwargs.get('eps'), + ce_kwargs.get('alpha') + ) * ce_kwargs.get('scale') - loss_dict['loss_fore_branch_triplet'] = triplet_loss( + if 'TripletLoss' in loss_names: + tri_kwargs = self.loss_kwargs.get('tri') + loss_dict['loss_fore_triplet'] = triplet_loss( fore_feat, gt_labels, - self._cfg.MODEL.LOSSES.TRI.MARGIN, - self._cfg.MODEL.LOSSES.TRI.NORM_FEAT, - self._cfg.MODEL.LOSSES.TRI.HARD_MINING, - ) * self._cfg.MODEL.LOSSES.TRI.SCALE - return loss_dict + tri_kwargs.get('margin'), + tri_kwargs.get('norm_feat'), + tri_kwargs.get('hard_mining') + ) * tri_kwargs.get('scale') + return loss_dict