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

update fastreid v1.2 

Summary:
1. refactor dataloader and heads
2. bugfix in fastattr, fastclas, fastface and partialreid
3. partial-fc supported in fastface
pull/456/head
liaoxingyu 2021-04-02 21:33:13 +08:00
parent 9288db6303
commit 44cee30dfc
40 changed files with 864 additions and 480 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
configs/VeRi/sbs_R50-ibn.yml
View File

@ -11,7 +11,6 @@ MODEL:
SOLVER:
OPT: SGD
NESTEROV: True
BASE_LR: 0.01
ETA_MIN_LR: 7.7e-5

4
fastreid/config/defaults.py
View File

@ -73,8 +73,8 @@ _C.MODEL.HEADS.POOL_LAYER = "GlobalAvgPool"
_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

41
fastreid/data/build.py
View File

@ -26,21 +26,19 @@ __all__ = [
_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
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)
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 = 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,
}

4
fastreid/engine/defaults.py
View File

@ -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 @@ class DefaultTrainer(TrainerBase):
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):

1
fastreid/engine/hooks.py
View File

@ -360,6 +360,7 @@ class EvalHook(HookBase):
)
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()

117
fastreid/layers/any_softmax.py
View File

@ -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
self.easy_margin = False
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 CosSoftmax(Linear):
r"""Implement of large margin cosine distance:
"""
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
def extra_repr(self):
return 'num_classes={}, scale={}, margin={}'.format(self._num_classes, self.s, self.m)
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
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 ArcSoftmax(Linear):
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.acos_()
logits[index] += m_hot
logits.cos_().mul_(self.s)
return logits
class CircleSoftmax(Linear):
def forward(self, logits, targets):
alpha_p = torch.clamp_min(-logits.detach() + 1 + self.m, min=0.)
@ -76,38 +62,19 @@ class CircleSoftmax(nn.Module):
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

10
fastreid/modeling/heads/embedding_head.py
View File

@ -83,15 +83,7 @@ class EmbeddingHead(nn.Module):
# 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

4
projects/FastAttr/fastattr/__init__.py
View File

@ -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

3
projects/FastAttr/fastattr/config.py
View File

@ -10,8 +10,7 @@ from fastreid.config import CfgNode as CN
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

74
projects/FastAttr/fastattr/data_build.py
View File

@ -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

9
projects/FastAttr/fastattr/modeling/__init__.py 100644
View File

@ -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

0
projects/FastAttr/fastattr/attr_baseline.py → projects/FastAttr/fastattr/modeling/attr_baseline.py
View File

0
projects/FastAttr/fastattr/attr_head.py → projects/FastAttr/fastattr/modeling/attr_head.py
View File

0
projects/FastAttr/fastattr/bce_loss.py → projects/FastAttr/fastattr/modeling/bce_loss.py
View File

49
projects/FastAttr/train_net.py
View File

@ -3,6 +3,7 @@
@author: xingyu liao
@contact: sherlockliao01@gmail.com
"""
import logging
import sys
sys.path.append('.')
@ -11,11 +12,15 @@ from fastreid.config import get_cfg
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 @@ class Trainer(DefaultTrainer):
"""
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()

23
projects/FastClas/configs/base-clas.yaml
View File

@ -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
OUTPUT_DIR: projects/FastClas/logs/r18_demo

2
projects/FastClas/fastclas/dataset.py
View File

@ -12,7 +12,7 @@ from fastreid.data.data_utils import read_image
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

39
projects/FastClas/train_net.py
View File

@ -7,8 +7,8 @@
import json
import logging
import sys
import os
import sys
sys.path.append('.')
@ -19,11 +19,14 @@ from fastreid.evaluation.clas_evaluator import ClasEvaluator
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 @@ class Trainer(DefaultTrainer):
"""
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 @@ class Trainer(DefaultTrainer):
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()

5
projects/FastFace/README.md
View File

@ -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 |
| ir100_se | 99.65 | 99.69 | 97.10 | 97.98 | 96.00 | 91.53 | 94.62 |
| ir50_se_0.1 | | | | | | | |
| ir100_se_0.1 | | | | | | | |

41
projects/FastFace/configs/face_base.yml
View File

@ -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

5
projects/FastFace/configs/r101_ir.yml
View File

@ -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

7
projects/FastFace/configs/r50_ir.yml
View File

@ -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

7
projects/FastFace/fastface/__init__.py
View File

@ -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

46
projects/FastFace/fastface/build.py
View File

@ -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

16
projects/FastFace/fastface/config.py 100644
View File

@ -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

2
projects/FastFace/fastface/datasets/ms1mv2.py
View File

@ -21,11 +21,9 @@ class MS1MV2(ImageDataset):
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):

80
projects/FastFace/fastface/face_data.py 100644
View File

@ -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,
}

10
projects/FastFace/fastface/modeling/__init__.py 100644
View File

@ -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

24
projects/FastFace/fastface/modeling/face_baseline.py 100644
View File

@ -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

39
projects/FastFace/fastface/modeling/face_head.py 100644
View File

@ -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,
}

196
projects/FastFace/fastface/modeling/partial_fc.py 100644
View File

@ -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

0
projects/FastFace/fastface/resnet_ir.py → projects/FastFace/fastface/modeling/resnet_ir.py
View File

173
projects/FastFace/fastface/trainer.py 100644
View File

@ -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()

31
projects/FastFace/train_net.py
View File

@ -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()

72
projects/PartialReID/configs/partial_market.yml
View File

@ -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"

4
projects/PartialReID/partialreid/config.py
View File

@ -10,6 +10,4 @@ from fastreid.config import CfgNode as CN
def add_partialreid_config(cfg):
_C = cfg
_C.TEST.DSR = CN()
_C.TEST.DSR.ENABLED = True
_C.TEST.DSR = CN({"ENABLED": True})

4
projects/PartialReID/partialreid/dsr_distance.py
View File

@ -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()

42
projects/PartialReID/partialreid/dsr_evaluation.py
View File

@ -10,11 +10,9 @@ from collections import OrderedDict
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 @@ class DsrEvaluator(DatasetEvaluator):
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 @@ class DsrEvaluator(DatasetEvaluator):
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)

98
projects/PartialReID/partialreid/dsr_head.py
View File

@ -9,8 +9,9 @@ import torch.nn.functional as F
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 @@ class OcclusionUnit(nn.Module):
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 @@ class OcclusionUnit(nn.Module):
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 @@ class OcclusionUnit(nn.Module):
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)
if with_bnneck:
occ_neck.append(get_norm(norm_type, feat_dim, bias_freeze=True))
self.bnneck_occ = nn.Sequential(*occ_neck)
self.bnneck_occ.apply(weights_init_kaiming)
# 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.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 @@ class DSRHead(nn.Module):
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 @@ class DSRHead(nn.Module):
# 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],
}

61
projects/PartialReID/partialreid/partialbaseline.py
View File

@ -12,58 +12,35 @@ from fastreid.modeling.meta_arch.build import META_ARCH_REGISTRY
@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
ce_kwargs.get('eps'),
ce_kwargs.get('alpha')
) * ce_kwargs.get('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
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