Repository: bupt-ai-cz/HHCL-ReID Branch: main Commit: 4986b46251de Files: 41 Total size: 128.1 KB Directory structure: gitextract_d7xkgx8j/ ├── README.md ├── examples/ │ ├── test.py │ └── train.py ├── hhcl/ │ ├── __init__.py │ ├── datasets/ │ │ ├── __init__.py │ │ ├── celebreid.py │ │ ├── dukemtmcreid.py │ │ ├── market1501.py │ │ ├── msmt17.py │ │ └── personx.py │ ├── evaluation_metrics/ │ │ ├── __init__.py │ │ ├── classification.py │ │ └── ranking.py │ ├── evaluators.py │ ├── models/ │ │ ├── __init__.py │ │ ├── cm.py │ │ ├── dsbn.py │ │ ├── kmeans.py │ │ ├── losses.py │ │ ├── pooling.py │ │ ├── resnet.py │ │ ├── resnet_ibn.py │ │ ├── resnet_ibn_a.py │ │ └── triplet.py │ ├── trainers.py │ └── utils/ │ ├── __init__.py │ ├── data/ │ │ ├── __init__.py │ │ ├── base_dataset.py │ │ ├── preprocessor.py │ │ ├── sampler.py │ │ └── transforms.py │ ├── faiss_rerank.py │ ├── faiss_utils.py │ ├── logging.py │ ├── meters.py │ ├── osutils.py │ ├── rerank.py │ └── serialization.py ├── requirements.txt ├── run.sh └── setup.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: README.md ================================================ # HHCL-ReID ![visitors](https://visitor-badge.glitch.me/badge?page_id=bupt-ai-cz.HHCL-ReID) [![Tweet](https://img.shields.io/twitter/url/http/shields.io.svg?style=social)](https://twitter.com/intent/tweet?text=Codes%20for%20Our%20Paper:%20"Hard-sample%20Guided%20Hybrid%20Contrast%20Learning%20for%20Unsupervised%20PersonRe-Identification"%20&url=https://github.com/bupt-ai-cz/HHCL-ReID) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/hard-sample-guided-hybrid-contrast-learning/unsupervised-person-re-identification-on-5)](https://paperswithcode.com/sota/unsupervised-person-re-identification-on-5?p=hard-sample-guided-hybrid-contrast-learning) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/hard-sample-guided-hybrid-contrast-learning/unsupervised-person-re-identification-on-4)](https://paperswithcode.com/sota/unsupervised-person-re-identification-on-4?p=hard-sample-guided-hybrid-contrast-learning) This repository is the official implementation of our paper "[Hard-sample Guided Hybrid Contrast Learning for Unsupervised Person Re-Identification](https://arxiv.org/abs/2109.12333)!". ![framework_HCCL](img/framework_HCCL.jpg) ## Requirements --- git clone https://github.com/bupt-ai-cz/HHCL-ReID.git cd HHCL-ReID pip install -r requirements.txt python setup.py develop ## Prepare Datasets --- Download the datasets Market-1501,MSMT17,DukeMTMC-reID from this [link](https://drive.google.com/file/d/19oWiYGjTgouFMK_psZvH8ysDGQ1KUbk-/view?usp=sharing) and unzip them under the directory like: HHCL-ReID/examples/data ├── market1501 │ └── Market-1501-v15.09.15 └── dukemtmcreid └── DukeMTMC-reID Prepare ImageNet Pre-trained Models for IBN-Net When training with the backbone of [IBN-ResNet](https://arxiv.org/abs/1807.09441), you need to download the ImageNet-pretrained model from this [link](https://drive.google.com/drive/folders/1thS2B8UOSBi_cJX6zRy6YYRwz_nVFI_S) and save it under the path of `examples/pretrained/`. ``` HHCL-ReID/examples └── pretrained └── resnet50_ibn_a.pth.tar ``` ## Training --- We utilize 4 GTX-2080TI GPUs for training. Examples: Market-1501: CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet50 -d market1501 --iters 200 --eps 0.45 --momentum 0.1 --num-instances 16 --pooling-type avg --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/market1501/resnet50_avg_cmhybrid DukeMTMC-reID: CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet50 -d dukemtmcreid --iters 200 --eps 0.6 --momentum 0.1 --num-instances 16 --pooling-type avg --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/dukemtmcreid/resnet50_avg_cmhybrid - use `-a resnet50` (default) for the backbone of ResNet-50, and `-a resnet_ibn50a` for the backbone of IBN-ResNet; - use `--pooling-type gem` for Generalized Mean Pooling (GEM) pooling and `--smooth` for label smoothing. ## Evaluation --- To evaluate my model on ImageNet, run: CUDA_VISIBLE_DEVICES=0 python examples/test.py -d $DATASET --resume $PATH --pooling-type avg ## Results --- Our model achieves the following performance on : | Dataset | Market1501 | | | | DukeMTMC-reID | | | | | ------------------ | ---------- | ---- | ---- | ---- | ------------- | ---- | ---- | ---- | | Setting | mAP | R1 | R5 | R10 | mAP | R1 | R5 | R10 | | Fully Unsupervised | 84.2 | 93.4 | 97.7 | 98.5 | 73.3 | 85.1 | 92.4 | 94.6 | | Supervised | 87.2 | 94.6 | 98.5 | 99.1 | 80.0 | 89.8 | 95.2 | 96.7 | You can download the above models in the paper from [Google Drive](https://drive.google.com/drive/folders/1WQw7wD2Mu_1SKl07_NdKvrYf2xrs3CEZ) ## Citation --- If you find this code useful for your research, please cite our paper ``` @article{hu2021hard, title={Hard-sample Guided Hybrid Contrast Learning for Unsupervised Person Re-Identification}, author={Hu, Zheng and Zhu, Chuang and He, Gang}, journal={arXiv preprint arXiv:2109.12333}, year={2021} } ``` ## Acknowledgements --- This project is not possible without multiple great opensourced codebases. We list them below. - [SpCL](https://github.com/yxgeee/SpCL) - [cluster-contrast-reid](https://github.com/alibaba/cluster-contrast-reid) ================================================ FILE: examples/test.py ================================================ from __future__ import print_function, absolute_import import argparse import os.path as osp import random import numpy as np import sys import torch from torch import nn from torch.backends import cudnn from torch.utils.data import DataLoader from hhcl import datasets from hhcl import models from hhcl.models.dsbn import convert_dsbn, convert_bn from hhcl.evaluators import Evaluator from hhcl.utils.data import transforms as T from hhcl.utils.data.preprocessor import Preprocessor from hhcl.utils.logging import Logger from hhcl.utils.serialization import load_checkpoint, save_checkpoint, copy_state_dict def get_data(name, data_dir, height, width, batch_size, workers): root = osp.join(data_dir, name) dataset = datasets.create(name, root) normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) test_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.ToTensor(), normalizer ]) test_loader = DataLoader( Preprocessor(list(set(dataset.query) | set(dataset.gallery)), root=dataset.images_dir, transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) return dataset, test_loader def main(): args = parser.parse_args() if args.seed is not None: random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.deterministic = True main_worker(args) def main_worker(args): cudnn.benchmark = True log_dir = osp.dirname(args.resume) sys.stdout = Logger(osp.join(log_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) # Create data loaders dataset, test_loader = get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # Create model model = models.create(args.arch, pretrained=False, num_features=args.features, dropout=args.dropout, num_classes=0, pooling_type=args.pooling_type) if args.dsbn: print("==> Load the model with domain-specific BNs") convert_dsbn(model) # Load from checkpoint checkpoint = load_checkpoint(args.resume) copy_state_dict(checkpoint['state_dict'], model, strip='module.') if args.dsbn: print("==> Test with {}-domain BNs".format("source" if args.test_source else "target")) convert_bn(model, use_target=(not args.test_source)) model.cuda() model = nn.DataParallel(model) # Evaluator model.eval() evaluator = Evaluator(model) evaluator.evaluate(test_loader, dataset.query, dataset.gallery, cmc_flag=True, rerank=args.rerank) return if __name__ == '__main__': parser = argparse.ArgumentParser(description="Testing the model") # data parser.add_argument('-d', '--dataset', type=str, default='market1501') parser.add_argument('-b', '--batch-size', type=int, default=256) parser.add_argument('-j', '--workers', type=int, default=4) parser.add_argument('--height', type=int, default=256, help="input height") parser.add_argument('--width', type=int, default=128, help="input width") # model parser.add_argument('-a', '--arch', type=str, default='resnet50', choices=models.names()) parser.add_argument('--features', type=int, default=0) parser.add_argument('--dropout', type=float, default=0) parser.add_argument('--resume', type=str, default="examples/logs/market1501/resnet50_avg/model_best.pth.tar", metavar='PATH') # testing configs parser.add_argument('--rerank', action='store_true', help="evaluation only") parser.add_argument('--dsbn', action='store_true', help="test on the model with domain-specific BN") parser.add_argument('--test-source', action='store_true', help="test on the source domain") parser.add_argument('--seed', type=int, default=1) # path working_dir = osp.dirname(osp.abspath(__file__)) parser.add_argument('--data-dir', type=str, metavar='PATH', default='examples/data') parser.add_argument('--pooling-type', type=str, default='avg') parser.add_argument('--embedding_features_path', type=str, default='examples/logs/market1501/resnet50_avg/') main() ================================================ FILE: examples/train.py ================================================ # -*- coding: utf-8 -*- from __future__ import print_function, absolute_import import argparse import os.path as osp import random import numpy as np import sys import collections import time from datetime import timedelta from sklearn.cluster import DBSCAN import torch from torch import nn from torch.backends import cudnn from torch.utils.data import DataLoader import torch.nn.functional as F from hhcl import datasets from hhcl import models from hhcl.models.cm import ClusterMemory from hhcl.trainers import Trainer from hhcl.evaluators import Evaluator, extract_features from hhcl.utils.data import IterLoader from hhcl.utils.data import transforms as T from hhcl.utils.data.sampler import RandomMultipleGallerySampler from hhcl.utils.data.preprocessor import Preprocessor from hhcl.utils.logging import Logger from hhcl.utils.serialization import load_checkpoint, save_checkpoint, copy_state_dict from hhcl.utils.faiss_rerank import compute_jaccard_distance start_epoch = best_mAP = 0 def get_data(name, data_dir): root = osp.join(data_dir, name) dataset = datasets.create(name, root) return dataset def get_train_loader(args, dataset, height, width, batch_size, workers, num_instances, iters, trainset=None): normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.RandomHorizontalFlip(p=0.5), T.Pad(10), T.RandomCrop((height, width)), T.ToTensor(), normalizer, T.RandomErasing(probability=0.5, mean=[0.485, 0.456, 0.406]) ]) train_set = sorted(dataset.train) if trainset is None else sorted(trainset) rmgs_flag = num_instances > 0 if rmgs_flag: sampler = RandomMultipleGallerySampler(train_set, num_instances) else: sampler = None train_loader = IterLoader( DataLoader(Preprocessor(train_set, root=dataset.images_dir, transform=train_transformer), batch_size=batch_size, num_workers=workers, sampler=sampler, shuffle=not rmgs_flag, pin_memory=True, drop_last=True), length=iters) return train_loader def get_test_loader(dataset, height, width, batch_size, workers, testset=None): normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) test_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.ToTensor(), normalizer ]) if testset is None: testset = list(set(dataset.query) | set(dataset.gallery)) test_loader = DataLoader( Preprocessor(testset, root=dataset.images_dir, transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) return test_loader def create_model(args): model = models.create(args.arch, num_features=args.features, norm=True, dropout=args.dropout, num_classes=0, pooling_type=args.pooling_type) # Load from checkpoint if args.resume: global start_epoch checkpoint = load_checkpoint(args.resume) copy_state_dict(checkpoint['state_dict'], model, strip='module.') start_epoch = checkpoint['epoch'] # use CUDA model.cuda() model = nn.DataParallel(model) return model def main(): args = parser.parse_args() if args.seed is not None: random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.deterministic = True main_worker(args) def main_worker(args): global start_epoch, best_mAP start_time = time.monotonic() cudnn.benchmark = True sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) print("==========\nArgs:{}\n==========".format(args)) # Create datasets iters = args.iters if (args.iters > 0) else None print("==> Load unlabeled dataset") dataset = get_data(args.dataset, args.data_dir) test_loader = get_test_loader(dataset, args.height, args.width, args.batch_size, args.workers) # Create model model = create_model(args) # Evaluator evaluator = Evaluator(model) # Optimizer params = [{"params": [value]} for _, value in model.named_parameters() if value.requires_grad] optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay) lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=0.1) # Trainer trainer = Trainer(model) for epoch in range(start_epoch, args.epochs): with torch.no_grad(): print('==> Create pseudo labels for unlabeled data') cluster_loader = get_test_loader(dataset, args.height, args.width, args.batch_size, args.workers, testset=sorted(dataset.train)) features, _ = extract_features(model, cluster_loader, print_freq=50) features = torch.cat([features[f].unsqueeze(0) for f, _, _ in sorted(dataset.train)], 0) rerank_dist = compute_jaccard_distance(features, k1=args.k1, k2=args.k2) if epoch == start_epoch: # DBSCAN cluster eps = args.eps print('Clustering criterion eps: {:.3f}'.format(eps)) cluster = DBSCAN(eps=eps, min_samples=4, metric='precomputed', n_jobs=-1) # select & cluster images as training set of this epochs pseudo_labels = cluster.fit_predict(rerank_dist) num_cluster = len(set(pseudo_labels)) - (1 if -1 in pseudo_labels else 0) # generate new dataset and calculate cluster centers @torch.no_grad() def generate_cluster_features(labels, features): centers = collections.defaultdict(list) for i, label in enumerate(labels): if label == -1: continue centers[labels[i]].append(features[i]) centers = [ torch.stack(centers[idx], dim=0).mean(0) for idx in sorted(centers.keys()) ] centers = torch.stack(centers, dim=0) return centers cluster_features = generate_cluster_features(pseudo_labels, features) def generate_random_features(labels, features, num_cluster, num_instances): indexes = np.zeros(num_cluster*num_instances) for i in range(num_cluster): index = [i+k*num_cluster for k in range(num_instances)] samples = np.random.choice(np.where(pseudo_labels==i)[0], num_instances, True) indexes[index] = samples memory_features = features[indexes] return memory_features if args.memorybank=='CMhybrid_v2': memory_features = generate_random_features(pseudo_labels, features, num_cluster, args.num_instances) mask = (pseudo_labels < 0).astype(int) print('==> Statistics for outliers with pseudo labels. outliers/total = {}/{} = {:.3f}'.format(mask.sum(), pseudo_labels.size, mask.sum()/pseudo_labels.size)) del cluster_loader, features # Create memory bank memory = ClusterMemory(model.module.num_features, num_cluster, temp=args.temp, momentum=args.momentum, mode=args.memorybank, smooth=args.smooth, num_instances=args.num_instances).cuda() if args.memorybank=='CMhybrid': memory.features = F.normalize(cluster_features.repeat(2, 1), dim=1).cuda() elif args.memorybank=='CMhybrid_v2': memory.features = F.normalize(torch.cat([cluster_features, memory_features],dim=0), dim=1).cuda() else: memory.features = F.normalize(cluster_features, dim=1).cuda() trainer.memory = memory pseudo_labeled_dataset = [] for i, ((fname, _, cid), label) in enumerate(zip(sorted(dataset.train), pseudo_labels)): if label != -1: pseudo_labeled_dataset.append((fname, label.item(), cid)) print('==> Statistics for epoch {}: {} clusters'.format(epoch, num_cluster)) train_loader = get_train_loader(args, dataset, args.height, args.width, args.batch_size, args.workers, args.num_instances, iters, trainset=pseudo_labeled_dataset) train_loader.new_epoch() trainer.train(epoch, train_loader, optimizer, print_freq=args.print_freq, train_iters=len(train_loader)) if (epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1): mAP = evaluator.evaluate(test_loader, dataset.query, dataset.gallery, cmc_flag=False) is_best = (mAP > best_mAP) best_mAP = max(mAP, best_mAP) save_checkpoint({ 'state_dict': model.state_dict(), 'epoch': epoch + 1, 'best_mAP': best_mAP, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} model mAP: {:5.1%} best: {:5.1%}{}\n'. format(epoch, mAP, best_mAP, ' *' if is_best else '')) lr_scheduler.step() print('==> Test with the best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model.load_state_dict(checkpoint['state_dict']) evaluator.evaluate(test_loader, dataset.query, dataset.gallery, cmc_flag=True) end_time = time.monotonic() print('Total running time: ', timedelta(seconds=end_time - start_time)) if __name__ == '__main__': parser = argparse.ArgumentParser(description="Hard-sample Guided Hybrid Contrast Learning for Unsupervised Person Re-ID") # data parser.add_argument('-d', '--dataset', type=str, default='dukemtmcreid', choices=datasets.names()) parser.add_argument('-b', '--batch-size', type=int, default=2) parser.add_argument('-j', '--workers', type=int, default=4) parser.add_argument('--height', type=int, default=256, help="input height") parser.add_argument('--width', type=int, default=128, help="input width") parser.add_argument('--num-instances', type=int, default=4, help="each minibatch consist of " "(batch_size // num_instances) identities, and " "each identity has num_instances instances, " "default: 0 (NOT USE)") # cluster parser.add_argument('--eps', type=float, default=0.6, help="max neighbor distance for DBSCAN") parser.add_argument('--eps-gap', type=float, default=0.02, help="multi-scale criterion for measuring cluster reliability") parser.add_argument('--k1', type=int, default=30, help="hyperparameter for jaccard distance") parser.add_argument('--k2', type=int, default=6, help="hyperparameter for jaccard distance") # model parser.add_argument('-a', '--arch', type=str, default='resnet50', choices=models.names()) parser.add_argument('--features', type=int, default=0) parser.add_argument('--dropout', type=float, default=0) parser.add_argument('--smooth', type=float, default=0, help="label smoothing") parser.add_argument('--hard-weight', type=float, default=0.5, help="hard weights") parser.add_argument('--momentum', type=float, default=0.1, help="update momentum for the memory bank") parser.add_argument('--pooling-type', type=str, default='gem') parser.add_argument('-mb', '--memorybank', type=str, default='CM', choices=['CM', 'CMhard', 'CMhybrid', 'CMhybrid_v2']) # optimizer parser.add_argument('--lr', type=float, default=0.00035, help="learning rate") parser.add_argument('--weight-decay', type=float, default=5e-4) parser.add_argument('--epochs', type=int, default=50) parser.add_argument('--iters', type=int, default=400) parser.add_argument('--step-size', type=int, default=20) # training configs parser.add_argument('--seed', type=int, default=1) parser.add_argument('--print-freq', type=int, default=10) parser.add_argument('--eval-step', type=int, default=10) parser.add_argument('--temp', type=float, default=0.05, help="temperature for scaling contrastive loss") # path working_dir = osp.dirname(osp.abspath(__file__)) parser.add_argument('--data-dir', type=str, metavar='PATH', default=osp.join(working_dir, 'data')) parser.add_argument('--logs-dir', type=str, metavar='PATH', default=osp.join(working_dir, 'logs')) parser.add_argument('--resume', type=str, metavar='PATH', default='') main() ================================================ FILE: hhcl/__init__.py ================================================ from __future__ import absolute_import from . import datasets from . import evaluation_metrics from . import models from . import utils from . import evaluators from . import trainers __version__ = '0.1.0' ================================================ FILE: hhcl/datasets/__init__.py ================================================ from __future__ import absolute_import import warnings from .market1501 import Market1501 from .msmt17 import MSMT17 from .personx import PersonX from .dukemtmcreid import DukeMTMCreID from .celebreid import CelebReID __factory = { 'market1501': Market1501, 'msmt17': MSMT17, 'personx': PersonX, 'dukemtmcreid': DukeMTMCreID, 'celebreid': CelebReID } def names(): return sorted(__factory.keys()) def create(name, root, *args, **kwargs): """ Create a dataset instance. Parameters ---------- name : str The dataset name. root : str The path to the dataset directory. split_id : int, optional The index of data split. Default: 0 num_val : int or float, optional When int, it means the number of validation identities. When float, it means the proportion of validation to all the trainval. Default: 100 download : bool, optional If True, will download the dataset. Default: False """ if name not in __factory: raise KeyError("Unknown dataset:", name) return __factory[name](root, *args, **kwargs) def get_dataset(name, root, *args, **kwargs): warnings.warn("get_dataset is deprecated. Use create instead.") return create(name, root, *args, **kwargs) ================================================ FILE: hhcl/datasets/celebreid.py ================================================ from __future__ import print_function, absolute_import import os.path as osp import glob import re from ..utils.data import BaseImageDataset class CelebReID(BaseImageDataset): """ CelebReID """ dataset_dir = 'CelebReID' def __init__(self, root, verbose=True, **kwargs): super(CelebReID, self).__init__() self.dataset_dir = osp.join(root, self.dataset_dir) self.train_dir = osp.join(self.dataset_dir, 'train') self.query_dir = osp.join(self.dataset_dir, 'query') self.gallery_dir = osp.join(self.dataset_dir, 'gallery') self._check_before_run() train = self._process_dir(self.train_dir, relabel=True) query = self._process_dir(self.query_dir, relabel=False) gallery = self._process_dir(self.gallery_dir, relabel=False) if verbose: print("=> CelebReID loaded") self.print_dataset_statistics(train, query, gallery) self.train = train self.query = query self.gallery = gallery self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train) self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query) self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery) def _check_before_run(self): """Check if all files are available before going deeper""" if not osp.exists(self.dataset_dir): raise RuntimeError("'{}' is not available".format(self.dataset_dir)) if not osp.exists(self.train_dir): raise RuntimeError("'{}' is not available".format(self.train_dir)) if not osp.exists(self.query_dir): raise RuntimeError("'{}' is not available".format(self.query_dir)) if not osp.exists(self.gallery_dir): raise RuntimeError("'{}' is not available".format(self.gallery_dir)) def _process_dir(self, dir_path, relabel=False): img_paths = glob.glob(osp.join(dir_path, '*.jpg')) # pattern = re.compile(r'([-\d]+)_c(\d)') pattern = re.compile(r'([-\d]+)_(\d)') pid_container = set() for img_path in img_paths: pid, _ = map(int, pattern.search(img_path).groups()) if pid == -1: continue # junk images are just ignored pid_container.add(pid) pid2label = {pid: label for label, pid in enumerate(pid_container)} dataset = [] for img_path in img_paths: pid, camid = map(int, pattern.search(img_path).groups()) if pid == -1: continue # junk images are just ignored # assert 0 <= pid <= 1501 # pid == 0 means background # assert 1 <= camid <= 6 camid -= 1 # index starts from 0 if relabel: pid = pid2label[pid] dataset.append((img_path, pid, camid)) return dataset ================================================ FILE: hhcl/datasets/dukemtmcreid.py ================================================ import glob import os.path as osp import re from ..utils.data import BaseImageDataset def process_dir(dir_path, relabel=False): img_paths = glob.glob(osp.join(dir_path, "*.jpg")) pattern = re.compile(r"([-\d]+)_c(\d)") # get all identities pid_container = set() for img_path in img_paths: pid, _ = map(int, pattern.search(img_path).groups()) if pid == -1: continue pid_container.add(pid) pid2label = {pid: label for label, pid in enumerate(pid_container)} data = [] for img_path in img_paths: pid, camid = map(int, pattern.search(img_path).groups()) if (pid not in pid_container) or (pid == -1): continue assert 1 <= camid <= 8 camid -= 1 if relabel: pid = pid2label[pid] data.append((img_path, pid, camid)) return data class DukeMTMCreID(BaseImageDataset): """DukeMTMC-reID. Reference: - Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016. - Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017. URL: ``_ Dataset statistics: - identities: 1404 (train + query). - images:16522 (train) + 2228 (query) + 17661 (gallery). - cameras: 8. """ dataset_dir = "DukeMTMC-reID" def __init__(self, root, verbose=True): super(DukeMTMCreID, self).__init__() self.root = osp.abspath(osp.expanduser(root)) self.dataset_dir = osp.join(self.root, self.dataset_dir) self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train') self.query_dir = osp.join(self.dataset_dir, 'query') self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test') train = process_dir(dir_path=self.train_dir, relabel=True) query = process_dir(dir_path=self.query_dir, relabel=False) gallery = process_dir(dir_path=self.gallery_dir, relabel=False) self.train = train self.query = query self.gallery = gallery self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train) self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query) self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery) def _check_before_run(self): """Check if all files are available before going deeper""" if not osp.exists(self.dataset_dir): raise RuntimeError("'{}' is not available".format(self.dataset_dir)) if not osp.exists(self.train_dir): raise RuntimeError("'{}' is not available".format(self.train_dir)) if not osp.exists(self.query_dir): raise RuntimeError("'{}' is not available".format(self.query_dir)) if not osp.exists(self.gallery_dir): raise RuntimeError("'{}' is not available".format(self.gallery_dir)) ================================================ FILE: hhcl/datasets/market1501.py ================================================ from __future__ import print_function, absolute_import import os.path as osp import glob import re from ..utils.data import BaseImageDataset class Market1501(BaseImageDataset): """ Market1501 Reference: Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015. URL: http://www.liangzheng.org/Project/project_reid.html Dataset statistics: # identities: 1501 (+1 for background) # images: 12936 (train) + 3368 (query) + 15913 (gallery) """ dataset_dir = 'Market-1501-v15.09.15' def __init__(self, root, verbose=True, **kwargs): super(Market1501, self).__init__() self.dataset_dir = osp.join(root, self.dataset_dir) self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train') self.query_dir = osp.join(self.dataset_dir, 'query') self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test') self._check_before_run() train = self._process_dir(self.train_dir, relabel=True) query = self._process_dir(self.query_dir, relabel=False) gallery = self._process_dir(self.gallery_dir, relabel=False) if verbose: print("=> Market1501 loaded") self.print_dataset_statistics(train, query, gallery) self.train = train self.query = query self.gallery = gallery self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train) self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query) self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery) def _check_before_run(self): """Check if all files are available before going deeper""" if not osp.exists(self.dataset_dir): raise RuntimeError("'{}' is not available".format(self.dataset_dir)) if not osp.exists(self.train_dir): raise RuntimeError("'{}' is not available".format(self.train_dir)) if not osp.exists(self.query_dir): raise RuntimeError("'{}' is not available".format(self.query_dir)) if not osp.exists(self.gallery_dir): raise RuntimeError("'{}' is not available".format(self.gallery_dir)) def _process_dir(self, dir_path, relabel=False): img_paths = glob.glob(osp.join(dir_path, '*.jpg')) pattern = re.compile(r'([-\d]+)_c(\d)') pid_container = set() for img_path in img_paths: pid, _ = map(int, pattern.search(img_path).groups()) if pid == -1: continue # junk images are just ignored pid_container.add(pid) pid2label = {pid: label for label, pid in enumerate(pid_container)} dataset = [] for img_path in img_paths: pid, camid = map(int, pattern.search(img_path).groups()) if pid == -1: continue # junk images are just ignored assert 0 <= pid <= 1501 # pid == 0 means background assert 1 <= camid <= 6 camid -= 1 # index starts from 0 if relabel: pid = pid2label[pid] dataset.append((img_path, pid, camid)) return dataset ================================================ FILE: hhcl/datasets/msmt17.py ================================================ from __future__ import print_function, absolute_import import os.path as osp import glob import re from ..utils.data import BaseImageDataset def _process_dir(dir_path, relabel=False): img_paths = glob.glob(osp.join(dir_path, '*.jpg')) pattern = re.compile(r'([-\d]+)_c(\d+)') pid_container = set() for img_path in img_paths: pid, _ = map(int, pattern.search(img_path).groups()) if pid == -1: continue # junk images are just ignored pid_container.add(pid) pid2label = {pid: label for label, pid in enumerate(pid_container)} dataset = [] for img_path in img_paths: pid, camid = map(int, pattern.search(img_path).groups()) if pid == -1: continue # junk images are just ignored assert 1 <= camid <= 15 camid -= 1 # index starts from 0 if relabel: pid = pid2label[pid] dataset.append((img_path, pid, camid)) return dataset class MSMT17(BaseImageDataset): dataset_dir = 'MSMT17_V1' def __init__(self, root, verbose=True, **kwargs): super(MSMT17, self).__init__() self.dataset_dir = osp.join(root, self.dataset_dir) self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train') self.query_dir = osp.join(self.dataset_dir, 'query') self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test') self._check_before_run() train = _process_dir(self.train_dir, relabel=True) query = _process_dir(self.query_dir, relabel=False) gallery = _process_dir(self.gallery_dir, relabel=False) if verbose: print("=> MSMT17_V1 loaded") self.print_dataset_statistics(train, query, gallery) self.train = train self.query = query self.gallery = gallery self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train) self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query) self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery) def _check_before_run(self): """Check if all files are available before going deeper""" if not osp.exists(self.dataset_dir): raise RuntimeError("'{}' is not available".format(self.dataset_dir)) if not osp.exists(self.train_dir): raise RuntimeError("'{}' is not available".format(self.train_dir)) if not osp.exists(self.query_dir): raise RuntimeError("'{}' is not available".format(self.query_dir)) if not osp.exists(self.gallery_dir): raise RuntimeError("'{}' is not available".format(self.gallery_dir)) ================================================ FILE: hhcl/datasets/personx.py ================================================ from __future__ import print_function, absolute_import import os.path as osp import glob import re from ..utils.data import BaseImageDataset class PersonX(BaseImageDataset): """ PersonX Reference: Sun et al. Dissecting Person Re-identification from the Viewpoint of Viewpoint. CVPR 2019. Dataset statistics: # identities: 1266 # images: 9840 (train) + 5136 (query) + 30816 (gallery) """ dataset_dir = 'PersonX' def __init__(self, root, verbose=True, **kwargs): super(PersonX, self).__init__() self.dataset_dir = osp.join(root, self.dataset_dir) self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train') self.query_dir = osp.join(self.dataset_dir, 'query') self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test') self._check_before_run() train = self._process_dir(self.train_dir, relabel=True) query = self._process_dir(self.query_dir, relabel=False) gallery = self._process_dir(self.gallery_dir, relabel=False) if verbose: print("=> PersonX loaded") self.print_dataset_statistics(train, query, gallery) self.train = train self.query = query self.gallery = gallery self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train) self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query) self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery) def _check_before_run(self): """Check if all files are available before going deeper""" if not osp.exists(self.dataset_dir): raise RuntimeError("'{}' is not available".format(self.dataset_dir)) if not osp.exists(self.train_dir): raise RuntimeError("'{}' is not available".format(self.train_dir)) if not osp.exists(self.query_dir): raise RuntimeError("'{}' is not available".format(self.query_dir)) if not osp.exists(self.gallery_dir): raise RuntimeError("'{}' is not available".format(self.gallery_dir)) def _process_dir(self, dir_path, relabel=False): img_paths = glob.glob(osp.join(dir_path, '*.jpg')) pattern = re.compile(r'([-\d]+)_c([-\d]+)') cam2label = {3: 1, 4: 2, 8: 3, 10: 4, 11: 5, 12: 6} pid_container = set() for img_path in img_paths: pid, _ = map(int, pattern.search(img_path).groups()) pid_container.add(pid) pid2label = {pid: label for label, pid in enumerate(pid_container)} dataset = [] for img_path in img_paths: pid, camid = map(int, pattern.search(img_path).groups()) assert (camid in cam2label.keys()) camid = cam2label[camid] camid -= 1 # index starts from 0 if relabel: pid = pid2label[pid] dataset.append((img_path, pid, camid)) return dataset ================================================ FILE: hhcl/evaluation_metrics/__init__.py ================================================ from __future__ import absolute_import from .classification import accuracy from .ranking import cmc, mean_ap __all__ = [ 'accuracy', 'cmc', 'mean_ap' ] ================================================ FILE: hhcl/evaluation_metrics/classification.py ================================================ from __future__ import absolute_import import torch from ..utils import to_torch def accuracy(output, target, topk=(1,)): with torch.no_grad(): output, target = to_torch(output), to_torch(target) maxk = max(topk) batch_size = target.size(0) _, pred = output.topk(maxk, 1, True, True) pred = pred.t() correct = pred.eq(target.view(1, -1).expand_as(pred)) ret = [] for k in topk: correct_k = correct[:k].view(-1).float().sum(dim=0, keepdim=True) ret.append(correct_k.mul_(1. / batch_size)) return ret ================================================ FILE: hhcl/evaluation_metrics/ranking.py ================================================ from __future__ import absolute_import from collections import defaultdict import numpy as np from sklearn.metrics import average_precision_score from ..utils import to_numpy def _unique_sample(ids_dict, num): mask = np.zeros(num, dtype=np.bool) for _, indices in ids_dict.items(): i = np.random.choice(indices) mask[i] = True return mask def cmc(distmat, query_ids=None, gallery_ids=None, query_cams=None, gallery_cams=None, topk=100, separate_camera_set=False, single_gallery_shot=False, first_match_break=False): distmat = to_numpy(distmat) m, n = distmat.shape # Fill up default values if query_ids is None: query_ids = np.arange(m) if gallery_ids is None: gallery_ids = np.arange(n) if query_cams is None: query_cams = np.zeros(m).astype(np.int32) if gallery_cams is None: gallery_cams = np.ones(n).astype(np.int32) # Ensure numpy array query_ids = np.asarray(query_ids) gallery_ids = np.asarray(gallery_ids) query_cams = np.asarray(query_cams) gallery_cams = np.asarray(gallery_cams) # Sort and find correct matches indices = np.argsort(distmat, axis=1) matches = (gallery_ids[indices] == query_ids[:, np.newaxis]) # Compute CMC for each query ret = np.zeros(topk) num_valid_queries = 0 for i in range(m): # Filter out the same id and same camera valid = ((gallery_ids[indices[i]] != query_ids[i]) | (gallery_cams[indices[i]] != query_cams[i])) if separate_camera_set: # Filter out samples from same camera valid &= (gallery_cams[indices[i]] != query_cams[i]) if not np.any(matches[i, valid]): continue if single_gallery_shot: repeat = 10 gids = gallery_ids[indices[i][valid]] inds = np.where(valid)[0] ids_dict = defaultdict(list) for j, x in zip(inds, gids): ids_dict[x].append(j) else: repeat = 1 for _ in range(repeat): if single_gallery_shot: # Randomly choose one instance for each id sampled = (valid & _unique_sample(ids_dict, len(valid))) index = np.nonzero(matches[i, sampled])[0] else: index = np.nonzero(matches[i, valid])[0] delta = 1. / (len(index) * repeat) for j, k in enumerate(index): if k - j >= topk: break if first_match_break: ret[k - j] += 1 break ret[k - j] += delta num_valid_queries += 1 if num_valid_queries == 0: raise RuntimeError("No valid query") return ret.cumsum() / num_valid_queries def mean_ap(distmat, query_ids=None, gallery_ids=None, query_cams=None, gallery_cams=None): distmat = to_numpy(distmat) m, n = distmat.shape # Fill up default values if query_ids is None: query_ids = np.arange(m) if gallery_ids is None: gallery_ids = np.arange(n) if query_cams is None: query_cams = np.zeros(m).astype(np.int32) if gallery_cams is None: gallery_cams = np.ones(n).astype(np.int32) # Ensure numpy array query_ids = np.asarray(query_ids) gallery_ids = np.asarray(gallery_ids) query_cams = np.asarray(query_cams) gallery_cams = np.asarray(gallery_cams) # Sort and find correct matches indices = np.argsort(distmat, axis=1) matches = (gallery_ids[indices] == query_ids[:, np.newaxis]) # Compute AP for each query aps = [] for i in range(m): # Filter out the same id and same camera valid = ((gallery_ids[indices[i]] != query_ids[i]) | (gallery_cams[indices[i]] != query_cams[i])) y_true = matches[i, valid] y_score = -distmat[i][indices[i]][valid] if not np.any(y_true): continue aps.append(average_precision_score(y_true, y_score)) if len(aps) == 0: raise RuntimeError("No valid query") return np.mean(aps) ================================================ FILE: hhcl/evaluators.py ================================================ from __future__ import print_function, absolute_import import time import collections from collections import OrderedDict import numpy as np import torch import random import copy from .evaluation_metrics import cmc, mean_ap from .utils.meters import AverageMeter from .utils.rerank import re_ranking from .utils import to_torch def extract_cnn_feature(model, inputs): inputs = to_torch(inputs).cuda() outputs = model(inputs) outputs = outputs.data.cpu() return outputs def extract_features(model, data_loader, print_freq=50): model.eval() batch_time = AverageMeter() data_time = AverageMeter() features = OrderedDict() labels = OrderedDict() end = time.time() with torch.no_grad(): for i, (imgs, fnames, pids, _, _) in enumerate(data_loader): data_time.update(time.time() - end) outputs = extract_cnn_feature(model, imgs) for fname, output, pid in zip(fnames, outputs, pids): features[fname] = output labels[fname] = pid batch_time.update(time.time() - end) end = time.time() if (i + 1) % print_freq == 0: print('Extract Features: [{}/{}]\t' 'Time {:.3f} ({:.3f})\t' 'Data {:.3f} ({:.3f})\t' .format(i + 1, len(data_loader), batch_time.val, batch_time.avg, data_time.val, data_time.avg)) return features, labels def pairwise_distance(features, query=None, gallery=None): if query is None and gallery is None: n = len(features) x = torch.cat(list(features.values())) x = x.view(n, -1) dist_m = torch.pow(x, 2).sum(dim=1, keepdim=True) * 2 dist_m = dist_m.expand(n, n) - 2 * torch.mm(x, x.t()) return dist_m x = torch.cat([features[f].unsqueeze(0) for f, _, _ in query], 0) y = torch.cat([features[f].unsqueeze(0) for f, _, _ in gallery], 0) m, n = x.size(0), y.size(0) x = x.view(m, -1) y = y.view(n, -1) dist_m = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(m, n) + \ torch.pow(y, 2).sum(dim=1, keepdim=True).expand(n, m).t() dist_m.addmm_(1, -2, x, y.t()) return dist_m, x.numpy(), y.numpy() def evaluate_all(query_features, gallery_features, distmat, query=None, gallery=None, query_ids=None, gallery_ids=None, query_cams=None, gallery_cams=None, cmc_topk=(1, 5, 10), cmc_flag=False): if query is not None and gallery is not None: query_ids = [pid for _, pid, _ in query] gallery_ids = [pid for _, pid, _ in gallery] query_cams = [cam for _, _, cam in query] gallery_cams = [cam for _, _, cam in gallery] else: assert (query_ids is not None and gallery_ids is not None and query_cams is not None and gallery_cams is not None) # Compute mean AP mAP = mean_ap(distmat, query_ids, gallery_ids, query_cams, gallery_cams) print('Mean AP: {:4.1%}'.format(mAP)) cmc_configs = { 'market1501': dict(separate_camera_set=False, single_gallery_shot=False, first_match_break=True),} cmc_scores = {name: cmc(distmat, query_ids, gallery_ids, query_cams, gallery_cams, **params) for name, params in cmc_configs.items()} print('CMC Scores:') for k in cmc_topk: print(' top-{:<4}{:12.1%}'.format(k, cmc_scores['market1501'][k-1])) if (not cmc_flag): return mAP return cmc_scores['market1501'], mAP class Evaluator(object): def __init__(self, model): super(Evaluator, self).__init__() self.model = model def evaluate(self, data_loader, query, gallery, cmc_flag=False, rerank=False): features, _ = extract_features(self.model, data_loader) distmat, query_features, gallery_features = pairwise_distance(features, query, gallery) results = evaluate_all(query_features, gallery_features, distmat, query=query, gallery=gallery, cmc_flag=cmc_flag) if (not rerank): return results print('Applying person re-ranking ...') distmat_qq, _, _ = pairwise_distance(features, query, query) distmat_gg, _, _ = pairwise_distance(features, gallery, gallery) distmat = re_ranking(distmat.numpy(), distmat_qq.numpy(), distmat_gg.numpy()) return evaluate_all(query_features, gallery_features, distmat, query=query, gallery=gallery, cmc_flag=cmc_flag) ================================================ FILE: hhcl/models/__init__.py ================================================ from __future__ import absolute_import from .resnet import * from .resnet_ibn import * __factory = { 'resnet18': resnet18, 'resnet34': resnet34, 'resnet50': resnet50, 'resnet101': resnet101, 'resnet152': resnet152, 'resnet_ibn50a': resnet_ibn50a, 'resnet_ibn101a': resnet_ibn101a, } def names(): return sorted(__factory.keys()) def create(name, *args, **kwargs): """ Create a model instance. Parameters ---------- name : str Model name. Can be one of 'inception', 'resnet18', 'resnet34', 'resnet50', 'resnet101', and 'resnet152'. pretrained : bool, optional Only applied for 'resnet*' models. If True, will use ImageNet pretrained model. Default: True cut_at_pooling : bool, optional If True, will cut the model before the last global pooling layer and ignore the remaining kwargs. Default: False num_features : int, optional If positive, will append a Linear layer after the global pooling layer, with this number of output units, followed by a BatchNorm layer. Otherwise these layers will not be appended. Default: 256 for 'inception', 0 for 'resnet*' norm : bool, optional If True, will normalize the feature to be unit L2-norm for each sample. Otherwise will append a ReLU layer after the above Linear layer if num_features > 0. Default: False dropout : float, optional If positive, will append a Dropout layer with this dropout rate. Default: 0 num_classes : int, optional If positive, will append a Linear layer at the end as the classifier with this number of output units. Default: 0 """ if name not in __factory: raise KeyError("Unknown model:", name) return __factory[name](*args, **kwargs) ================================================ FILE: hhcl/models/cm.py ================================================ import collections import numpy as np from abc import ABC import torch import torch.nn.functional as F from torch import nn, autograd from .losses import CrossEntropyLabelSmooth, FocalTopLoss class CM(autograd.Function): @staticmethod def forward(ctx, inputs, targets, features, momentum): ctx.features = features ctx.momentum = momentum ctx.save_for_backward(inputs, targets) outputs = inputs.mm(ctx.features.t()) return outputs @staticmethod def backward(ctx, grad_outputs): inputs, targets = ctx.saved_tensors grad_inputs = None if ctx.needs_input_grad[0]: grad_inputs = grad_outputs.mm(ctx.features) # momentum update for x, y in zip(inputs, targets): ctx.features[y] = ctx.momentum * ctx.features[y] + (1. - ctx.momentum) * x ctx.features[y] /= ctx.features[y].norm() return grad_inputs, None, None, None def cm(inputs, indexes, features, momentum=0.5): return CM.apply(inputs, indexes, features, torch.Tensor([momentum]).to(inputs.device)) class CM_Hard(autograd.Function): @staticmethod def forward(ctx, inputs, targets, features, momentum): ctx.features = features ctx.momentum = momentum ctx.save_for_backward(inputs, targets) outputs = inputs.mm(ctx.features.t()) return outputs @staticmethod def backward(ctx, grad_outputs): inputs, targets = ctx.saved_tensors grad_inputs = None if ctx.needs_input_grad[0]: grad_inputs = grad_outputs.mm(ctx.features) batch_centers = collections.defaultdict(list) for instance_feature, index in zip(inputs, targets.tolist()): batch_centers[index].append(instance_feature) for index, features in batch_centers.items(): distances = [] for feature in features: distance = feature.unsqueeze(0).mm(ctx.features[index].unsqueeze(0).t())[0][0] distances.append(distance.cpu().numpy()) median = np.argmin(np.array(distances)) ctx.features[index] = ctx.features[index] * ctx.momentum + (1 - ctx.momentum) * features[median] ctx.features[index] /= ctx.features[index].norm() return grad_inputs, None, None, None def cm_hard(inputs, indexes, features, momentum=0.5): return CM_Hard.apply(inputs, indexes, features, torch.Tensor([momentum]).to(inputs.device)) class CM_Hybrid(autograd.Function): @staticmethod def forward(ctx, inputs, targets, features, momentum): ctx.features = features ctx.momentum = momentum ctx.save_for_backward(inputs, targets) outputs = inputs.mm(ctx.features.t()) return outputs @staticmethod def backward(ctx, grad_outputs): inputs, targets = ctx.saved_tensors nums = len(ctx.features)//2 grad_inputs = None if ctx.needs_input_grad[0]: grad_inputs = grad_outputs.mm(ctx.features) batch_centers = collections.defaultdict(list) for instance_feature, index in zip(inputs, targets.tolist()): batch_centers[index].append(instance_feature) for index, features in batch_centers.items(): distances = [] for feature in features: distance = feature.unsqueeze(0).mm(ctx.features[index].unsqueeze(0).t())[0][0] distances.append(distance.cpu().numpy()) median = np.argmin(np.array(distances)) ctx.features[index] = ctx.features[index] * ctx.momentum + (1 - ctx.momentum) * features[median] ctx.features[index] /= ctx.features[index].norm() mean = torch.stack(features, dim=0).mean(0) ctx.features[index+nums] = ctx.features[index+nums] * ctx.momentum + (1 - ctx.momentum) * mean ctx.features[index+nums] /= ctx.features[index+nums].norm() return grad_inputs, None, None, None def cm_hybrid(inputs, indexes, features, momentum=0.5): return CM_Hybrid.apply(inputs, indexes, features, torch.Tensor([momentum]).to(inputs.device)) class CM_Hybrid_v2(autograd.Function): @staticmethod def forward(ctx, inputs, targets, features, momentum, num_instances): ctx.features = features ctx.momentum = momentum ctx.num_instances = num_instances ctx.save_for_backward(inputs, targets) outputs = inputs.mm(ctx.features.t()) return outputs @staticmethod def backward(ctx, grad_outputs): inputs, targets = ctx.saved_tensors nums = len(ctx.features)//(ctx.num_instances + 1) grad_inputs = None if ctx.needs_input_grad[0]: grad_inputs = grad_outputs.mm(ctx.features) batch_centers = collections.defaultdict(list) updated = set() for k, (instance_feature, index) in enumerate(zip(inputs, targets.tolist())): batch_centers[index].append(instance_feature) if index not in updated: indexes = [index + nums*i for i in range(1, (targets==index).sum()+1)] ctx.features[indexes] = inputs[targets==index] # ctx.features[indexes] = ctx.features[indexes] * ctx.momentum + (1 - ctx.momentum) * inputs[targets==index] # ctx.features[indexes] /= ctx.features[indexes].norm(dim=1, keepdim=True) updated.add(index) for index, features in batch_centers.items(): mean = torch.stack(features, dim=0).mean(0) ctx.features[index] = ctx.features[index] * ctx.momentum + (1 - ctx.momentum) * mean ctx.features[index] /= ctx.features[index].norm() return grad_inputs, None, None, None, None def cm_hybrid_v2(inputs, indexes, features, momentum=0.5, num_instances=16, *args): return CM_Hybrid_v2.apply(inputs, indexes, features, torch.Tensor([momentum]).to(inputs.device), num_instances) class ClusterMemory(nn.Module, ABC): __CMfactory = { 'CM': cm, 'CMhard':cm_hard, } def __init__(self, num_features, num_samples, temp=0.05, momentum=0.2, mode='CM', hard_weight=0.5, smooth=0., num_instances=1): super(ClusterMemory, self).__init__() self.num_features = num_features self.num_samples = num_samples self.momentum = momentum self.temp = temp self.cm_type = mode if smooth>0: self.cross_entropy = CrossEntropyLabelSmooth(self.num_samples, 0.1, True) print('>>> Using CrossEntropy with Label Smoothing.') else: self.cross_entropy = nn.CrossEntropyLoss().cuda() if self.cm_type in ['CM', 'CMhard']: self.register_buffer('features', torch.zeros(num_samples, num_features)) elif self.cm_type=='CMhybrid': self.hard_weight = hard_weight print('hard_weight: {}'.format(self.hard_weight)) self.register_buffer('features', torch.zeros(2*num_samples, num_features)) elif self.cm_type=='CMhybrid_v2': self.hard_weight = hard_weight self.num_instances = num_instances self.register_buffer('features', torch.zeros((self.num_instances+1)*num_samples, num_features)) else: raise TypeError('Cluster Memory {} is invalid!'.format(self.cm_type)) def forward(self, inputs, targets): if self.cm_type in ['CM', 'CMhard']: outputs = ClusterMemory.__CMfactory[self.cm_type](inputs, targets, self.features, self.momentum) outputs /= self.temp loss = self.cross_entropy(outputs, targets) return loss elif self.cm_type=='CMhybrid': outputs = cm_hybrid(inputs, targets, self.features, self.momentum) outputs /= self.temp output_hard, output_mean = torch.chunk(outputs, 2, dim=1) loss = self.hard_weight * (self.cross_entropy(output_hard, targets) + (1 - self.hard_weight) * self.cross_entropy(output_mean, targets)) return loss elif self.cm_type=='CMhybrid_v2': outputs = cm_hybrid_v2(inputs, targets, self.features, self.momentum, self.num_instances) out_list = torch.chunk(outputs, self.num_instances+1, dim=1) out = torch.stack(out_list[1:], dim=0) neg = torch.max(out, dim=0)[0] pos = torch.min(out, dim=0)[0] mask = torch.zeros_like(out_list[0]).scatter_(1, targets.unsqueeze(1), 1) logits = mask * pos + (1-mask) * neg loss = self.hard_weight * self.cross_entropy(out_list[0]/self.temp, targets) \ + (1 - self.hard_weight) * self.cross_entropy(logits/self.temp, targets) return loss ================================================ FILE: hhcl/models/dsbn.py ================================================ import torch import torch.nn as nn # Domain-specific BatchNorm class DSBN2d(nn.Module): def __init__(self, planes): super(DSBN2d, self).__init__() self.num_features = planes self.BN_S = nn.BatchNorm2d(planes) self.BN_T = nn.BatchNorm2d(planes) def forward(self, x): if (not self.training): return self.BN_T(x) bs = x.size(0) assert (bs%2==0) split = torch.split(x, int(bs/2), 0) out1 = self.BN_S(split[0].contiguous()) out2 = self.BN_T(split[1].contiguous()) out = torch.cat((out1, out2), 0) return out class DSBN1d(nn.Module): def __init__(self, planes): super(DSBN1d, self).__init__() self.num_features = planes self.BN_S = nn.BatchNorm1d(planes) self.BN_T = nn.BatchNorm1d(planes) def forward(self, x): if (not self.training): return self.BN_T(x) bs = x.size(0) assert (bs%2==0) split = torch.split(x, int(bs/2), 0) out1 = self.BN_S(split[0].contiguous()) out2 = self.BN_T(split[1].contiguous()) out = torch.cat((out1, out2), 0) return out def convert_dsbn(model): for _, (child_name, child) in enumerate(model.named_children()): assert(not next(model.parameters()).is_cuda) if isinstance(child, nn.BatchNorm2d): m = DSBN2d(child.num_features) m.BN_S.load_state_dict(child.state_dict()) m.BN_T.load_state_dict(child.state_dict()) setattr(model, child_name, m) elif isinstance(child, nn.BatchNorm1d): m = DSBN1d(child.num_features) m.BN_S.load_state_dict(child.state_dict()) m.BN_T.load_state_dict(child.state_dict()) setattr(model, child_name, m) else: convert_dsbn(child) def convert_bn(model, use_target=True): for _, (child_name, child) in enumerate(model.named_children()): assert(not next(model.parameters()).is_cuda) if isinstance(child, DSBN2d): m = nn.BatchNorm2d(child.num_features) if use_target: m.load_state_dict(child.BN_T.state_dict()) else: m.load_state_dict(child.BN_S.state_dict()) setattr(model, child_name, m) elif isinstance(child, DSBN1d): m = nn.BatchNorm1d(child.num_features) if use_target: m.load_state_dict(child.BN_T.state_dict()) else: m.load_state_dict(child.BN_S.state_dict()) setattr(model, child_name, m) else: convert_bn(child, use_target=use_target) ================================================ FILE: hhcl/models/kmeans.py ================================================ # Written by Yixiao Ge import warnings import faiss import torch from ..utils import to_numpy, to_torch __all__ = ["label_generator_kmeans"] @torch.no_grad() def label_generator_kmeans(features, num_classes=500, cuda=True): assert num_classes, "num_classes for kmeans is null" # k-means cluster by faiss cluster = faiss.Kmeans( features.size(-1), num_classes, niter=300, verbose=True, gpu=cuda ) cluster.train(to_numpy(features)) _, labels = cluster.index.search(to_numpy(features), 1) labels = labels.reshape(-1) centers = to_torch(cluster.centroids).float() # labels = to_torch(labels).long() # k-means does not have outlier points assert not (-1 in labels) return labels, centers, num_classes, None ================================================ FILE: hhcl/models/losses.py ================================================ import torch import torch.nn as nn import torch.nn.functional as F from torch.nn import Parameter class CrossEntropyLabelSmooth(nn.Module): """Cross entropy loss with label smoothing regularizer. Reference: Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016. Equation: y = (1 - epsilon) * y + epsilon / K. Args: num_classes (int): number of classes. epsilon (float): weight. """ def __init__(self, num_classes=0, epsilon=0.1, topk_smoothing=False): super(CrossEntropyLabelSmooth, self).__init__() self.num_classes = num_classes self.epsilon = epsilon self.logsoftmax = nn.LogSoftmax(dim=1).cuda() self.k = 1 if not topk_smoothing else self.num_classes//50 def forward(self, inputs, targets): """ Args: inputs: prediction matrix (before softmax) with shape (batch_size, num_classes) targets: ground truth labels with shape (num_classes) """ log_probs = self.logsoftmax(inputs) if self.k >1: topk = torch.argsort(-log_probs)[:,:self.k] targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1 - self.epsilon) targets += torch.zeros_like(log_probs).scatter_(1, topk, self.epsilon / self.k) else: targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1) targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes loss = (- targets * log_probs).mean(0).sum() return loss class SoftEntropy(nn.Module): def __init__(self, input_prob=False): super(SoftEntropy, self).__init__() self.input_prob = input_prob self.logsoftmax = nn.LogSoftmax(dim=1).cuda() def forward(self, inputs, targets): log_probs = self.logsoftmax(inputs) if self.input_prob: loss = (- targets.detach() * log_probs).mean(0).sum() else: loss = (- F.softmax(targets, dim=1).detach() * log_probs).mean(0).sum() return loss class SoftEntropySmooth(nn.Module): def __init__(self, epsilon=0.1): super(SoftEntropySmooth, self).__init__() self.epsilon = epsilon self.logsoftmax = nn.LogSoftmax(dim=1).cuda() def forward(self, inputs, soft_targets, targets): log_probs = self.logsoftmax(inputs) targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1) soft_targets = F.softmax(soft_targets, dim=1) smooth_targets = (1 - self.epsilon) * targets + self.epsilon * soft_targets loss = (- smooth_targets.detach() * log_probs).mean(0).sum() return loss class Softmax(nn.Module): def __init__(self, feat_dim, num_class, temp=0.05): super(Softmax, self).__init__() self.weight = Parameter(torch.Tensor(feat_dim, num_class)) self.weight.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5) self.temp = temp def forward(self, feats, labels): kernel_norm = F.normalize(self.weight, dim=0) feats = F.normalize(feats) outputs = feats.mm(kernel_norm) outputs /= self.temp loss = F.cross_entropy(outputs, labels) return loss class CircleLoss(nn.Module): """Implementation for "Circle Loss: A Unified Perspective of Pair Similarity Optimization" Note: this is the classification based implementation of circle loss. """ def __init__(self, feat_dim, num_class, margin=0.25, gamma=256): super(CircleLoss, self).__init__() self.weight = Parameter(torch.Tensor(feat_dim, num_class)) self.weight.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5) self.margin = margin self.gamma = gamma self.O_p = 1 + margin self.O_n = -margin self.delta_p = 1-margin self.delta_n = margin def forward(self, feats, labels): kernel_norm = F.normalize(self.weight, dim=0) feats = F.normalize(feats) cos_theta = torch.mm(feats, kernel_norm) cos_theta = cos_theta.clamp(-1, 1) index_pos = torch.zeros_like(cos_theta) index_pos.scatter_(1, labels.data.view(-1, 1), 1) index_pos = index_pos.bool() index_neg = torch.ones_like(cos_theta) index_neg.scatter_(1, labels.data.view(-1, 1), 0) index_neg = index_neg.bool() alpha_p = torch.clamp_min(self.O_p - cos_theta.detach(), min=0.) alpha_n = torch.clamp_min(cos_theta.detach() - self.O_n, min=0.) logit_p = alpha_p * (cos_theta - self.delta_p) logit_n = alpha_n * (cos_theta - self.delta_n) output = cos_theta * 1.0 output[index_pos] = logit_p[index_pos] output[index_neg] = logit_n[index_neg] output *= self.gamma return F.cross_entropy(output, labels) class CosFace(nn.Module): r"""Implement of CosFace (https://arxiv.org/pdf/1801.09414.pdf): Args: in_features: size of each input sample out_features: size of each output sample s: norm of input feature m: margin cos(theta)-m """ def __init__(self, feat_dim, num_class, s = 64.0, m = 0.35): super(CosFace, self).__init__() self.in_features = feat_dim self.out_features = num_class self.s = s self.m = m self.weight = Parameter(torch.FloatTensor(feat_dim, num_class)) nn.init.xavier_uniform_(self.weight) def forward(self, input, label): # --------------------------- cos(theta) & phi(theta) --------------------------- # cosine = F.linear(F.normalize(input), F.normalize(self.weight, dim=1)) cosine = torch.mm(F.normalize(input), F.normalize(self.weight, dim=0)) phi = cosine - self.m # --------------------------- convert label to one-hot --------------------------- one_hot = torch.zeros(cosine.size(), device = 'cuda') # one_hot = one_hot.cuda() if cosine.is_cuda else one_hot one_hot.scatter_(1, label.view(-1, 1).long(), 1) # -------------torch.where(out_i = {x_i if condition_i else y_i) ------------- output = (one_hot * phi) + ((1.0 - one_hot) * cosine) # you can use torch.where if your torch.__version__ is 0.4 output *= self.s return F.cross_entropy(output, label) def __repr__(self): return self.__class__.__name__ + '(' \ + 'in_features = ' + str(self.in_features) \ + ', out_features = ' + str(self.out_features) \ + ', s = ' + str(self.s) \ + ', m = ' + str(self.m) + ')' import math class InstanceLoss(nn.Module): def __init__(self, batch_size, temperature, device): super(InstanceLoss, self).__init__() self.batch_size = batch_size self.temperature = temperature self.device = device self.mask = self.mask_correlated_samples(batch_size) self.criterion = nn.CrossEntropyLoss(reduction="sum") def mask_correlated_samples(self, batch_size): N = 2 * batch_size mask = torch.ones((N, N)) mask = mask.fill_diagonal_(0) for i in range(batch_size): mask[i, batch_size + i] = 0 mask[batch_size + i, i] = 0 mask = mask.bool() return mask def forward(self, z_i, z_j): N = 2 * self.batch_size z = torch.cat((z_i, z_j), dim=0) sim = torch.matmul(z, z.T) / self.temperature sim_i_j = torch.diag(sim, self.batch_size) sim_j_i = torch.diag(sim, -self.batch_size) positive_samples = torch.cat((sim_i_j, sim_j_i), dim=0).reshape(N, 1) negative_samples = sim[self.mask].reshape(N, -1) labels = torch.zeros(N).to(positive_samples.device).long() logits = torch.cat((positive_samples, negative_samples), dim=1) loss = self.criterion(logits, labels) loss /= N return loss class ClusterLoss(nn.Module): def __init__(self, class_num, temperature, device): super(ClusterLoss, self).__init__() self.class_num = class_num self.temperature = temperature self.device = device self.mask = self.mask_correlated_clusters(class_num) self.criterion = nn.CrossEntropyLoss(reduction="sum") self.similarity_f = nn.CosineSimilarity(dim=2) def mask_correlated_clusters(self, class_num): N = 2 * class_num mask = torch.ones((N, N)) mask = mask.fill_diagonal_(0) for i in range(class_num): mask[i, class_num + i] = 0 mask[class_num + i, i] = 0 mask = mask.bool() return mask def forward(self, c_i, c_j): p_i = c_i.sum(0).view(-1) p_i /= p_i.sum() ne_i = math.log(p_i.size(0)) + (p_i * torch.log(p_i)).sum() p_j = c_j.sum(0).view(-1) p_j /= p_j.sum() ne_j = math.log(p_j.size(0)) + (p_j * torch.log(p_j)).sum() ne_loss = ne_i + ne_j c_i = c_i.t() c_j = c_j.t() N = 2 * self.class_num c = torch.cat((c_i, c_j), dim=0) sim = self.similarity_f(c.unsqueeze(1), c.unsqueeze(0)) / self.temperature sim_i_j = torch.diag(sim, self.class_num) sim_j_i = torch.diag(sim, -self.class_num) positive_clusters = torch.cat((sim_i_j, sim_j_i), dim=0).reshape(N, 1) negative_clusters = sim[self.mask].reshape(N, -1) labels = torch.zeros(N).to(positive_clusters.device).long() logits = torch.cat((positive_clusters, negative_clusters), dim=1) loss = self.criterion(logits, labels) loss /= N return loss + ne_loss class FocalLoss(nn.Module): def __init__(self, gamma=2, alpha=0.25): super(FocalLoss, self).__init__() self.alpha = alpha self.gamma = gamma print('Initializing FocalLoss for training: alpha={}, gamma={}'.format(self.alpha, self.gamma)) def forward(self, input, target): assert input.dim() == 2 assert not target.requires_grad target = target.squeeze(1) if target.dim() == 2 else target assert target.dim() == 1 logpt = F.log_softmax(input, dim=1) logpt_gt = logpt.gather(1,target.unsqueeze(1)) logpt_gt = logpt_gt.view(-1) pt_gt = logpt_gt.exp() assert logpt_gt.size() == pt_gt.size() loss = -self.alpha*(torch.pow((1-pt_gt), self.gamma))*logpt_gt return loss.mean() class LabelRefineLoss(nn.Module): def __init__(self, lambda1=0.0): super(LabelRefineLoss, self).__init__() self.lambda1 = lambda1 print('Initializing LabelRefineLoss for training: lambda1={}'.format(self.lambda1)) def forward(self, input, target): assert input.dim() == 2 assert not target.requires_grad target = target.squeeze(1) if target.dim() == 2 else target assert target.dim() == 1 logpt = F.log_softmax(input, dim=1) logpt_gt = logpt.gather(1,target.unsqueeze(1)) logpt_gt = logpt_gt.view(-1) logpt_pred,_ = torch.max(logpt,1) logpt_pred = logpt_pred.view(-1) assert logpt_gt.size() == logpt_pred.size() loss = - (1-self.lambda1)*logpt_gt - self.lambda1* logpt_pred return loss.mean() class FocalTopLoss(nn.Module): def __init__(self, top_percent=0.7): super(FocalTopLoss, self).__init__() self.top_percent = top_percent def masked_softmax_multi_focal(self, vec, targets=None, dim=1): exps = torch.exp(vec) one_hot_pos = F.one_hot(targets, num_classes=exps.shape[1]) one_hot_neg = one_hot_pos.new_ones(size=one_hot_pos.shape) one_hot_neg = one_hot_neg - one_hot_pos neg_exps = exps.new_zeros(size=exps.shape) neg_exps[one_hot_neg>0] = exps[one_hot_neg>0] ori_neg_exps = neg_exps neg_exps = neg_exps/neg_exps.sum(dim=1, keepdim=True) new_exps = exps.new_zeros(size=exps.shape) new_exps[one_hot_pos>0] = exps[one_hot_pos>0] sorted, indices = torch.sort(neg_exps, dim=1, descending=True) sorted_cum_sum = torch.cumsum(sorted, dim=1) sorted_cum_diff = (sorted_cum_sum - self.top_percent).abs() sorted_cum_min_indices = sorted_cum_diff.argmin(dim=1) min_values = sorted[torch.range(0, sorted.shape[0]-1).long(), sorted_cum_min_indices] min_values = min_values.unsqueeze(dim=-1) * ori_neg_exps.sum(dim=1, keepdim=True) ori_neg_exps[ori_neg_exps0] = ori_neg_exps[one_hot_neg>0] masked_sums = exps.sum(dim, keepdim=True) return new_exps / masked_sums def forward(self, input, target): masked_sim = self.masked_softmax_multi_focal(input, target) return F.nll_loss(torch.log(masked_sim + 1e-6), target) ================================================ FILE: hhcl/models/pooling.py ================================================ # Credit to https://github.com/JDAI-CV/fast-reid/blob/master/fastreid/layers/pooling.py from abc import ABC import torch import torch.nn.functional as F from torch import nn __all__ = [ "GeneralizedMeanPoolingPFpn", "GeneralizedMeanPoolingList", "GeneralizedMeanPoolingP", "AdaptiveAvgMaxPool2d", "FastGlobalAvgPool2d", "avg_pooling", "max_pooling", ] class GeneralizedMeanPoolingList(nn.Module, ABC): r"""Applies a 2D power-average adaptive pooling over an input signal composed of several input planes. The function computed is: :math:`f(X) = pow(sum(pow(X, p)), 1/p)` - At p = infinity, one gets Max Pooling - At p = 1, one gets Average Pooling The output is of size H x W, for any input size. The number of output features is equal to the number of input planes. Args: output_size: the target output size of the image of the form H x W. Can be a tuple (H, W) or a single H for a square image H x H H and W can be either a ``int``, or ``None`` which means the size will be the same as that of the input. """ def __init__(self, output_size=1, eps=1e-6): super(GeneralizedMeanPoolingList, self).__init__() self.output_size = output_size self.eps = eps def forward(self, x_list): outs = [] for x in x_list: x = x.clamp(min=self.eps) out = torch.nn.functional.adaptive_avg_pool2d(x, self.output_size) outs.append(out) return torch.stack(outs, -1).mean(-1) def __repr__(self): return ( self.__class__.__name__ + "(" + "output_size=" + str(self.output_size) + ")" ) class GeneralizedMeanPooling(nn.Module, ABC): r"""Applies a 2D power-average adaptive pooling over an input signal composed of several input planes. The function computed is: :math:`f(X) = pow(sum(pow(X, p)), 1/p)` - At p = infinity, one gets Max Pooling - At p = 1, one gets Average Pooling The output is of size H x W, for any input size. The number of output features is equal to the number of input planes. Args: output_size: the target output size of the image of the form H x W. Can be a tuple (H, W) or a single H for a square image H x H H and W can be either a ``int``, or ``None`` which means the size will be the same as that of the input. """ def __init__(self, norm, output_size=1, eps=1e-6): super(GeneralizedMeanPooling, self).__init__() assert norm > 0 self.p = float(norm) self.output_size = output_size self.eps = eps def forward(self, x): x = x.clamp(min=self.eps).pow(self.p) return torch.nn.functional.adaptive_avg_pool2d(x, self.output_size).pow( 1.0 / self.p ) def __repr__(self): return ( self.__class__.__name__ + "(" + str(self.p) + ", " + "output_size=" + str(self.output_size) + ")" ) class GeneralizedMeanPoolingP(GeneralizedMeanPooling, ABC): """ Same, but norm is trainable """ def __init__(self, norm=3, output_size=1, eps=1e-6): super(GeneralizedMeanPoolingP, self).__init__(norm, output_size, eps) self.p = nn.Parameter(torch.ones(1) * norm) class GeneralizedMeanPoolingFpn(nn.Module, ABC): r"""Applies a 2D power-average adaptive pooling over an input signal composed of several input planes. The function computed is: :math:`f(X) = pow(sum(pow(X, p)), 1/p)` - At p = infinity, one gets Max Pooling - At p = 1, one gets Average Pooling The output is of size H x W, for any input size. The number of output features is equal to the number of input planes. Args: output_size: the target output size of the image of the form H x W. Can be a tuple (H, W) or a single H for a square image H x H H and W can be either a ``int``, or ``None`` which means the size will be the same as that of the input. """ def __init__(self, norm, output_size=1, eps=1e-6): super(GeneralizedMeanPoolingFpn, self).__init__() assert norm > 0 self.p = float(norm) self.output_size = output_size self.eps = eps def forward(self, x_lists): outs = [] for x in x_lists: x = x.clamp(min=self.eps).pow(self.p) out = torch.nn.functional.adaptive_avg_pool2d(x, self.output_size).pow( 1.0 / self.p ) outs.append(out) return torch.cat(outs, 1) def __repr__(self): return ( self.__class__.__name__ + "(" + str(self.p) + ", " + "output_size=" + str(self.output_size) + ")" ) class GeneralizedMeanPoolingPFpn(GeneralizedMeanPoolingFpn, ABC): """ Same, but norm is trainable """ def __init__(self, norm=3, output_size=1, eps=1e-6): super(GeneralizedMeanPoolingPFpn, self).__init__(norm, output_size, eps) self.p = nn.Parameter(torch.ones(1) * norm) class AdaptiveAvgMaxPool2d(nn.Module, ABC): def __init__(self): super(AdaptiveAvgMaxPool2d, self).__init__() self.avgpool = FastGlobalAvgPool2d() def forward(self, x): x_avg = self.avgpool(x, self.output_size) x_max = F.adaptive_max_pool2d(x, 1) x = x_max + x_avg return x class FastGlobalAvgPool2d(nn.Module, ABC): def __init__(self, flatten=False): super(FastGlobalAvgPool2d, self).__init__() self.flatten = flatten def forward(self, x): if self.flatten: in_size = x.size() return x.view((in_size[0], in_size[1], -1)).mean(dim=2) else: return ( x.view(x.size(0), x.size(1), -1) .mean(-1) .view(x.size(0), x.size(1), 1, 1) ) def avg_pooling(): return nn.AdaptiveAvgPool2d(1) # return FastGlobalAvgPool2d() def max_pooling(): return nn.AdaptiveMaxPool2d(1) class Flatten(nn.Module): def forward(self, input): return input.view(input.size(0), -1) __pooling_factory = { "avg": avg_pooling, "max": max_pooling, "gem": GeneralizedMeanPoolingP, "gemFpn": GeneralizedMeanPoolingPFpn, "gemList": GeneralizedMeanPoolingList, "avg+max": AdaptiveAvgMaxPool2d, } def pooling_names(): return sorted(__pooling_factory.keys()) def build_pooling_layer(name): """ Create a pooling layer. Parameters ---------- name : str The backbone name. """ if name not in __pooling_factory: raise KeyError("Unknown pooling layer:", name) return __pooling_factory[name]() ================================================ FILE: hhcl/models/resnet.py ================================================ from __future__ import absolute_import from torch import nn from torch.nn import functional as F from torch.nn import init import torchvision import torch from .pooling import build_pooling_layer __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152'] class ResNet(nn.Module): __factory = { 18: torchvision.models.resnet18, 34: torchvision.models.resnet34, 50: torchvision.models.resnet50, 101: torchvision.models.resnet101, 152: torchvision.models.resnet152, } def __init__(self, depth, pretrained=True, cut_at_pooling=False, num_features=0, norm=False, dropout=0, num_classes=0, pooling_type='avg'): print('pooling_type: {}'.format(pooling_type)) super(ResNet, self).__init__() self.pretrained = pretrained self.depth = depth self.cut_at_pooling = cut_at_pooling # Construct base (pretrained) resnet if depth not in ResNet.__factory: raise KeyError("Unsupported depth:", depth) resnet = ResNet.__factory[depth](pretrained=pretrained) if self.depth >= 50: resnet.layer4[0].conv2.stride = (1, 1) resnet.layer4[0].downsample[0].stride = (1, 1) self.base = nn.Sequential( resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4) self.gap = build_pooling_layer(pooling_type) if not self.cut_at_pooling: self.num_features = num_features self.norm = norm self.dropout = dropout self.has_embedding = num_features > 0 self.num_classes = num_classes out_planes = resnet.fc.in_features # Append new layers if self.has_embedding: self.feat = nn.Linear(out_planes, self.num_features) self.feat_bn = nn.BatchNorm1d(self.num_features) init.kaiming_normal_(self.feat.weight, mode='fan_out') init.constant_(self.feat.bias, 0) else: # Change the num_features to CNN output channels self.num_features = out_planes self.feat_bn = nn.BatchNorm1d(self.num_features) self.feat_bn.bias.requires_grad_(False) if self.dropout > 0: self.drop = nn.Dropout(self.dropout) if self.num_classes > 0: self.classifier = nn.Linear(self.num_features, self.num_classes, bias=False) init.normal_(self.classifier.weight, std=0.001) init.constant_(self.feat_bn.weight, 1) init.constant_(self.feat_bn.bias, 0) if not pretrained: self.reset_params() def forward(self, x): bs = x.size(0) x = self.base(x) x = self.gap(x) x = x.view(x.size(0), -1) if self.cut_at_pooling: return x if self.has_embedding: bn_x = self.feat_bn(self.feat(x)) else: bn_x = self.feat_bn(x) if (self.training is False): bn_x = F.normalize(bn_x) return bn_x if self.norm: bn_x = F.normalize(bn_x) elif self.has_embedding: bn_x = F.relu(bn_x) if self.dropout > 0: bn_x = self.drop(bn_x) if self.num_classes > 0: prob = self.classifier(bn_x) else: return bn_x return prob, bn_x def reset_params(self): for m in self.modules(): if isinstance(m, nn.Conv2d): init.kaiming_normal_(m.weight, mode='fan_out') if m.bias is not None: init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): init.constant_(m.weight, 1) init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm1d): init.constant_(m.weight, 1) init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): init.normal_(m.weight, std=0.001) if m.bias is not None: init.constant_(m.bias, 0) def resnet18(**kwargs): return ResNet(18, **kwargs) def resnet34(**kwargs): return ResNet(34, **kwargs) def resnet50(**kwargs): return ResNet(50, **kwargs) def resnet101(**kwargs): return ResNet(101, **kwargs) def resnet152(**kwargs): return ResNet(152, **kwargs) ================================================ FILE: hhcl/models/resnet_ibn.py ================================================ from __future__ import absolute_import from torch import nn from torch.nn import functional as F from torch.nn import init import torchvision import torch from .pooling import build_pooling_layer from .resnet_ibn_a import resnet50_ibn_a, resnet101_ibn_a __all__ = ['ResNetIBN', 'resnet_ibn50a', 'resnet_ibn101a'] class ResNetIBN(nn.Module): __factory = { '50a': resnet50_ibn_a, '101a': resnet101_ibn_a } def __init__(self, depth, pretrained=True, cut_at_pooling=False, num_features=0, norm=False, dropout=0, num_classes=0, pooling_type='avg'): print('pooling_type: {}'.format(pooling_type)) super(ResNetIBN, self).__init__() self.depth = depth self.pretrained = pretrained self.cut_at_pooling = cut_at_pooling resnet = ResNetIBN.__factory[depth](pretrained=pretrained) resnet.layer4[0].conv2.stride = (1, 1) resnet.layer4[0].downsample[0].stride = (1, 1) self.base = nn.Sequential( resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4) self.gap = build_pooling_layer(pooling_type) if not self.cut_at_pooling: self.num_features = num_features self.norm = norm self.dropout = dropout self.has_embedding = num_features > 0 self.num_classes = num_classes out_planes = resnet.fc.in_features # Append new layers if self.has_embedding: self.feat = nn.Linear(out_planes, self.num_features) self.feat_bn = nn.BatchNorm1d(self.num_features) init.kaiming_normal_(self.feat.weight, mode='fan_out') init.constant_(self.feat.bias, 0) else: # Change the num_features to CNN output channels self.num_features = out_planes self.feat_bn = nn.BatchNorm1d(self.num_features) self.feat_bn.bias.requires_grad_(False) if self.dropout > 0: self.drop = nn.Dropout(self.dropout) if self.num_classes > 0: self.classifier = nn.Linear(self.num_features, self.num_classes, bias=False) init.normal_(self.classifier.weight, std=0.001) init.constant_(self.feat_bn.weight, 1) init.constant_(self.feat_bn.bias, 0) if not pretrained: self.reset_params() def forward(self, x): x = self.base(x) x = self.gap(x) x = x.view(x.size(0), -1) if self.cut_at_pooling: return x if self.has_embedding: bn_x = self.feat_bn(self.feat(x)) else: bn_x = self.feat_bn(x) if self.training is False: bn_x = F.normalize(bn_x) return bn_x if self.norm: bn_x = F.normalize(bn_x) elif self.has_embedding: bn_x = F.relu(bn_x) if self.dropout > 0: bn_x = self.drop(bn_x) if self.num_classes > 0: prob = self.classifier(bn_x) else: return bn_x return prob, bn_x def reset_params(self): for m in self.modules(): if isinstance(m, nn.Conv2d): init.kaiming_normal_(m.weight, mode='fan_out') if m.bias is not None: init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): init.constant_(m.weight, 1) init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm1d): init.constant_(m.weight, 1) init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): init.normal_(m.weight, std=0.001) if m.bias is not None: init.constant_(m.bias, 0) def resnet_ibn50a(**kwargs): return ResNetIBN('50a', **kwargs) def resnet_ibn101a(**kwargs): return ResNetIBN('101a', **kwargs) ================================================ FILE: hhcl/models/resnet_ibn_a.py ================================================ import torch import torch.nn as nn import math import torch.utils.model_zoo as model_zoo __all__ = ['ResNet', 'resnet50_ibn_a', 'resnet101_ibn_a'] model_urls = { 'ibn_resnet50a': './examples/pretrained/resnet50_ibn_a.pth.tar', 'ibn_resnet101a': './examples/pretrained/resnet101_ibn_a.pth.tar', } def conv3x3(in_planes, out_planes, stride=1): "3x3 convolution with padding" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None): super(BasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = nn.BatchNorm2d(planes) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out class IBN(nn.Module): def __init__(self, planes): super(IBN, self).__init__() half1 = int(planes/2) self.half = half1 half2 = planes - half1 self.IN = nn.InstanceNorm2d(half1, affine=True) self.BN = nn.BatchNorm2d(half2) def forward(self, x): split = torch.split(x, self.half, 1) out1 = self.IN(split[0].contiguous()) out2 = self.BN(split[1].contiguous()) out = torch.cat((out1, out2), 1) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None): super(Bottleneck, self).__init__() self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) if ibn: self.bn1 = IBN(planes) else: self.bn1 = nn.BatchNorm2d(planes) self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(planes) self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False) self.bn3 = nn.BatchNorm2d(planes * self.expansion) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out class ResNet(nn.Module): def __init__(self, block, layers, num_classes=1000): scale = 64 self.inplanes = scale super(ResNet, self).__init__() self.conv1 = nn.Conv2d(3, scale, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = nn.BatchNorm2d(scale) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, scale, layers[0]) self.layer2 = self._make_layer(block, scale*2, layers[1], stride=2) self.layer3 = self._make_layer(block, scale*4, layers[2], stride=2) self.layer4 = self._make_layer(block, scale*8, layers[3], stride=2) self.avgpool = nn.AvgPool2d(7) self.fc = nn.Linear(scale * 8 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, nn.Conv2d): n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels m.weight.data.normal_(0, math.sqrt(2. / n)) elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() elif isinstance(m, nn.InstanceNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() def _make_layer(self, block, planes, blocks, stride=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(planes * block.expansion), ) layers = [] ibn = True if planes == 512: ibn = False layers.append(block(self.inplanes, planes, ibn, stride, downsample)) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes, ibn)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = x.view(x.size(0), -1) x = self.fc(x) return x def resnet50_ibn_a(pretrained=False, **kwargs): """Constructs a ResNet-50 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) if pretrained: state_dict = torch.load(model_urls['ibn_resnet50a'], map_location=torch.device('cpu'))['state_dict'] state_dict = remove_module_key(state_dict) model.load_state_dict(state_dict) return model def resnet101_ibn_a(pretrained=False, **kwargs): """Constructs a ResNet-101 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) if pretrained: state_dict = torch.load(model_urls['ibn_resnet101a'], map_location=torch.device('cpu'))['state_dict'] state_dict = remove_module_key(state_dict) model.load_state_dict(state_dict) return model def remove_module_key(state_dict): for key in list(state_dict.keys()): if 'module' in key: state_dict[key.replace('module.','')] = state_dict.pop(key) return state_dict ================================================ FILE: hhcl/models/triplet.py ================================================ from __future__ import absolute_import import torch from torch import nn import torch.nn.functional as F def euclidean_dist(x, y): m, n = x.size(0), y.size(0) xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n) yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t() dist = xx + yy dist.addmm_(1, -2, x, y.t()) dist = dist.clamp(min=1e-12).sqrt() # for numerical stability return dist def cosine_dist(x, y): bs1, bs2 = x.size(0), y.size(0) frac_up = torch.matmul(x, y.transpose(0, 1)) frac_down = (torch.sqrt(torch.sum(torch.pow(x, 2), 1))).view(bs1, 1).repeat(1, bs2) * \ (torch.sqrt(torch.sum(torch.pow(y, 2), 1))).view(1, bs2).repeat(bs1, 1) cosine = frac_up / frac_down return 1-cosine def _batch_hard(mat_distance, mat_similarity, indice=False): sorted_mat_distance, positive_indices = torch.sort(mat_distance + (-9999999.) * (1 - mat_similarity), dim=1, descending=True) hard_p = sorted_mat_distance[:, 0] hard_p_indice = positive_indices[:, 0] sorted_mat_distance, negative_indices = torch.sort(mat_distance + (9999999.) * (mat_similarity), dim=1, descending=False) hard_n = sorted_mat_distance[:, 0] hard_n_indice = negative_indices[:, 0] if(indice): return hard_p, hard_n, hard_p_indice, hard_n_indice return hard_p, hard_n class TripletLoss(nn.Module): ''' Compute Triplet loss augmented with Batch Hard Details can be seen in 'In defense of the Triplet Loss for Person Re-Identification' ''' def __init__(self, margin, normalize_feature=False): super(TripletLoss, self).__init__() self.margin = margin self.normalize_feature = normalize_feature self.margin_loss = nn.MarginRankingLoss(margin=margin).cuda() def forward(self, emb, label): if self.normalize_feature: # equal to cosine similarity emb = F.normalize(emb) mat_dist = euclidean_dist(emb, emb) # mat_dist = cosine_dist(emb, emb) assert mat_dist.size(0) == mat_dist.size(1) N = mat_dist.size(0) mat_sim = label.expand(N, N).eq(label.expand(N, N).t()).float() dist_ap, dist_an = _batch_hard(mat_dist, mat_sim) assert dist_an.size(0)==dist_ap.size(0) y = torch.ones_like(dist_ap) loss = self.margin_loss(dist_an, dist_ap, y) prec = (dist_an.data > dist_ap.data).sum() * 1. / y.size(0) return loss, prec class SoftTripletLoss(nn.Module): def __init__(self, margin=None, normalize_feature=False): super(SoftTripletLoss, self).__init__() self.margin = margin self.normalize_feature = normalize_feature def forward(self, emb1, emb2, label): if self.normalize_feature: # equal to cosine similarity emb1 = F.normalize(emb1) emb2 = F.normalize(emb2) mat_dist = euclidean_dist(emb1, emb1) assert mat_dist.size(0) == mat_dist.size(1) N = mat_dist.size(0) mat_sim = label.expand(N, N).eq(label.expand(N, N).t()).float() dist_ap, dist_an, ap_idx, an_idx = _batch_hard(mat_dist, mat_sim, indice=True) assert dist_an.size(0)==dist_ap.size(0) triple_dist = torch.stack((dist_ap, dist_an), dim=1) triple_dist = F.log_softmax(triple_dist, dim=1) if (self.margin is not None): loss = (- self.margin * triple_dist[:,0] - (1 - self.margin) * triple_dist[:,1]).mean() return loss mat_dist_ref = euclidean_dist(emb2, emb2) dist_ap_ref = torch.gather(mat_dist_ref, 1, ap_idx.view(N,1).expand(N,N))[:,0] dist_an_ref = torch.gather(mat_dist_ref, 1, an_idx.view(N,1).expand(N,N))[:,0] triple_dist_ref = torch.stack((dist_ap_ref, dist_an_ref), dim=1) triple_dist_ref = F.softmax(triple_dist_ref, dim=1).detach() loss = (- triple_dist_ref * triple_dist).mean(0).sum() return loss ================================================ FILE: hhcl/trainers.py ================================================ from __future__ import print_function, absolute_import import time import torch import torch.nn.functional as F from .utils.meters import AverageMeter class Trainer(object): def __init__(self, encoder, memory=None): super(Trainer, self).__init__() self.encoder = encoder self.memory = memory def train(self, epoch, data_loader, optimizer, print_freq=10, train_iters=400): self.encoder.train() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() end = time.time() for i in range(train_iters): # load data inputs = data_loader.next() data_time.update(time.time() - end) # process inputs inputs, labels, indexes = self._parse_data(inputs) loss = 0 # forward f_out = self._forward(inputs) loss += self.memory(f_out, labels) optimizer.zero_grad() loss.backward() optimizer.step() losses.update(loss.item()) # print log batch_time.update(time.time() - end) end = time.time() if (i + 1) % print_freq == 0: print('Epoch: [{}][{}/{}]\t' 'Time {:.3f} ({:.3f})\t' 'Data {:.3f} ({:.3f})\t' 'Loss {:.3f} ({:.3f})' .format(epoch, i + 1, len(data_loader), batch_time.val, batch_time.avg, data_time.val, data_time.avg, losses.val, losses.avg)) def _parse_data(self, inputs): imgs, _, pids, _, indexes = inputs return imgs.cuda(), pids.cuda(), indexes.cuda() def _forward(self, inputs): return self.encoder(inputs) ================================================ FILE: hhcl/utils/__init__.py ================================================ from __future__ import absolute_import import torch def to_numpy(tensor): if torch.is_tensor(tensor): return tensor.cpu().numpy() elif type(tensor).__module__ != 'numpy': raise ValueError("Cannot convert {} to numpy array" .format(type(tensor))) return tensor def to_torch(ndarray): if type(ndarray).__module__ == 'numpy': return torch.from_numpy(ndarray) elif not torch.is_tensor(ndarray): raise ValueError("Cannot convert {} to torch tensor" .format(type(ndarray))) return ndarray ================================================ FILE: hhcl/utils/data/__init__.py ================================================ from __future__ import absolute_import from .base_dataset import BaseDataset, BaseImageDataset from .preprocessor import Preprocessor class IterLoader: def __init__(self, loader, length=None): self.loader = loader self.length = length self.iter = None def __len__(self): if self.length is not None: return self.length return len(self.loader) def new_epoch(self): self.iter = iter(self.loader) def next(self): try: return next(self.iter) except: self.iter = iter(self.loader) return next(self.iter) ================================================ FILE: hhcl/utils/data/base_dataset.py ================================================ # encoding: utf-8 import numpy as np class BaseDataset(object): """ Base class of reid dataset """ def get_imagedata_info(self, data): pids, cams = [], [] for _, pid, camid in data: pids += [pid] cams += [camid] pids = set(pids) cams = set(cams) num_pids = len(pids) num_cams = len(cams) num_imgs = len(data) return num_pids, num_imgs, num_cams def print_dataset_statistics(self): raise NotImplementedError @property def images_dir(self): return None class BaseImageDataset(BaseDataset): """ Base class of image reid dataset """ def print_dataset_statistics(self, train, query, gallery): num_train_pids, num_train_imgs, num_train_cams = self.get_imagedata_info(train) num_query_pids, num_query_imgs, num_query_cams = self.get_imagedata_info(query) num_gallery_pids, num_gallery_imgs, num_gallery_cams = self.get_imagedata_info(gallery) print("Dataset statistics:") print(" ----------------------------------------") print(" subset | # ids | # images | # cameras") print(" ----------------------------------------") print(" train | {:5d} | {:8d} | {:9d}".format(num_train_pids, num_train_imgs, num_train_cams)) print(" query | {:5d} | {:8d} | {:9d}".format(num_query_pids, num_query_imgs, num_query_cams)) print(" gallery | {:5d} | {:8d} | {:9d}".format(num_gallery_pids, num_gallery_imgs, num_gallery_cams)) print(" ----------------------------------------") ================================================ FILE: hhcl/utils/data/preprocessor.py ================================================ from __future__ import absolute_import import os import os.path as osp from torch.utils.data import DataLoader, Dataset import numpy as np import random import math from PIL import Image class Preprocessor(Dataset): def __init__(self, dataset, root=None, transform=None, mutual=False): super(Preprocessor, self).__init__() self.dataset = dataset self.root = root self.transform = transform self.mutual = mutual def __len__(self): return len(self.dataset) def __getitem__(self, indices): if self.mutual: return self._get_mutual_item(indices) else: return self._get_single_item(indices) def _get_single_item(self, index): fname, pid, camid = self.dataset[index] fpath = fname if self.root is not None: fpath = osp.join(self.root, fname) img = Image.open(fpath).convert('RGB') if self.transform is not None: img = self.transform(img) return img, fname, pid, camid, index def _get_mutual_item(self, index): fname, pid, camid = self.dataset[index] fpath = fname if self.root is not None: fpath = osp.join(self.root, fname) img_1 = Image.open(fpath).convert('RGB') img_2 = img_1.copy() if self.transform is not None: img_1 = self.transform(img_1) img_2 = self.transform(img_2) return img_1, img_2, pid, camid ================================================ FILE: hhcl/utils/data/sampler.py ================================================ from __future__ import absolute_import from collections import defaultdict import math import numpy as np import copy import random import torch from torch.utils.data.sampler import ( Sampler, SequentialSampler, RandomSampler, SubsetRandomSampler, WeightedRandomSampler) def No_index(a, b): assert isinstance(a, list) return [i for i, j in enumerate(a) if j != b] class RandomIdentitySampler(Sampler): def __init__(self, data_source, num_instances): self.data_source = data_source self.num_instances = num_instances self.index_dic = defaultdict(list) for index, (_, pid, _) in enumerate(data_source): self.index_dic[pid].append(index) self.pids = list(self.index_dic.keys()) self.num_samples = len(self.pids) def __len__(self): return self.num_samples * self.num_instances def __iter__(self): indices = torch.randperm(self.num_samples).tolist() ret = [] for i in indices: pid = self.pids[i] t = self.index_dic[pid] if len(t) >= self.num_instances: t = np.random.choice(t, size=self.num_instances, replace=False) else: t = np.random.choice(t, size=self.num_instances, replace=True) ret.extend(t) return iter(ret) class RandomMultipleGallerySampler(Sampler): def __init__(self, data_source, num_instances=4): super().__init__(data_source) self.data_source = data_source self.index_pid = defaultdict(int) self.pid_cam = defaultdict(list) self.pid_index = defaultdict(list) self.num_instances = num_instances for index, (_, pid, cam) in enumerate(data_source): if pid < 0: continue self.index_pid[index] = pid self.pid_cam[pid].append(cam) self.pid_index[pid].append(index) self.pids = list(self.pid_index.keys()) self.num_samples = len(self.pids) def __len__(self): return self.num_samples * self.num_instances def __iter__(self): indices = torch.randperm(len(self.pids)).tolist() ret = [] for kid in indices: i = random.choice(self.pid_index[self.pids[kid]]) _, i_pid, i_cam = self.data_source[i] ret.append(i) pid_i = self.index_pid[i] cams = self.pid_cam[pid_i] index = self.pid_index[pid_i] select_cams = No_index(cams, i_cam) if select_cams: if len(select_cams) >= self.num_instances: cam_indexes = np.random.choice(select_cams, size=self.num_instances-1, replace=False) else: cam_indexes = np.random.choice(select_cams, size=self.num_instances-1, replace=True) for kk in cam_indexes: ret.append(index[kk]) else: select_indexes = No_index(index, i) if not select_indexes: continue if len(select_indexes) >= self.num_instances: ind_indexes = np.random.choice(select_indexes, size=self.num_instances-1, replace=False) else: ind_indexes = np.random.choice(select_indexes, size=self.num_instances-1, replace=True) for kk in ind_indexes: ret.append(index[kk]) return iter(ret) ================================================ FILE: hhcl/utils/data/transforms.py ================================================ from __future__ import absolute_import from torchvision.transforms import * from PIL import Image import random import math import numpy as np class RectScale(object): def __init__(self, height, width, interpolation=Image.BILINEAR): self.height = height self.width = width self.interpolation = interpolation def __call__(self, img): w, h = img.size if h == self.height and w == self.width: return img return img.resize((self.width, self.height), self.interpolation) class RandomSizedRectCrop(object): def __init__(self, height, width, interpolation=Image.BILINEAR): self.height = height self.width = width self.interpolation = interpolation def __call__(self, img): for attempt in range(10): area = img.size[0] * img.size[1] target_area = random.uniform(0.64, 1.0) * area aspect_ratio = random.uniform(2, 3) h = int(round(math.sqrt(target_area * aspect_ratio))) w = int(round(math.sqrt(target_area / aspect_ratio))) if w <= img.size[0] and h <= img.size[1]: x1 = random.randint(0, img.size[0] - w) y1 = random.randint(0, img.size[1] - h) img = img.crop((x1, y1, x1 + w, y1 + h)) assert(img.size == (w, h)) return img.resize((self.width, self.height), self.interpolation) # Fallback scale = RectScale(self.height, self.width, interpolation=self.interpolation) return scale(img) class RandomErasing(object): """ Randomly selects a rectangle region in an image and erases its pixels. 'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/pdf/1708.04896.pdf Args: probability: The probability that the Random Erasing operation will be performed. sl: Minimum proportion of erased area against input image. sh: Maximum proportion of erased area against input image. r1: Minimum aspect ratio of erased area. mean: Erasing value. """ def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4914, 0.4822, 0.4465)): self.probability = probability self.mean = mean self.sl = sl self.sh = sh self.r1 = r1 def __call__(self, img): if random.uniform(0, 1) >= self.probability: return img for attempt in range(100): area = img.size()[1] * img.size()[2] target_area = random.uniform(self.sl, self.sh) * area aspect_ratio = random.uniform(self.r1, 1 / self.r1) h = int(round(math.sqrt(target_area * aspect_ratio))) w = int(round(math.sqrt(target_area / aspect_ratio))) if w < img.size()[2] and h < img.size()[1]: x1 = random.randint(0, img.size()[1] - h) y1 = random.randint(0, img.size()[2] - w) if img.size()[0] == 3: img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] img[1, x1:x1 + h, y1:y1 + w] = self.mean[1] img[2, x1:x1 + h, y1:y1 + w] = self.mean[2] else: img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] return img return img ================================================ FILE: hhcl/utils/faiss_rerank.py ================================================ #!/usr/bin/env python3 # -*- coding: utf-8 -*- """ CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017. url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf Matlab version: https://github.com/zhunzhong07/person-re-ranking """ import os, sys import time import numpy as np from scipy.spatial.distance import cdist import gc import faiss import torch import torch.nn.functional as F from .faiss_utils import search_index_pytorch, search_raw_array_pytorch, \ index_init_gpu, index_init_cpu def k_reciprocal_neigh(initial_rank, i, k1): forward_k_neigh_index = initial_rank[i,:k1+1] backward_k_neigh_index = initial_rank[forward_k_neigh_index,:k1+1] fi = np.where(backward_k_neigh_index==i)[0] return forward_k_neigh_index[fi] def compute_jaccard_distance(target_features, k1=20, k2=6, print_flag=True, search_option=0, use_float16=False): end = time.time() if print_flag: print('Computing jaccard distance...') ngpus = faiss.get_num_gpus() N = target_features.size(0) mat_type = np.float16 if use_float16 else np.float32 if (search_option==0): # GPU + PyTorch CUDA Tensors (1) res = faiss.StandardGpuResources() res.setDefaultNullStreamAllDevices() _, initial_rank = search_raw_array_pytorch(res, target_features, target_features, k1) initial_rank = initial_rank.cpu().numpy() elif (search_option==1): # GPU + PyTorch CUDA Tensors (2) res = faiss.StandardGpuResources() index = faiss.GpuIndexFlatL2(res, target_features.size(-1)) index.add(target_features.cpu().numpy()) _, initial_rank = search_index_pytorch(index, target_features, k1) res.syncDefaultStreamCurrentDevice() initial_rank = initial_rank.cpu().numpy() elif (search_option==2): # GPU index = index_init_gpu(ngpus, target_features.size(-1)) index.add(target_features.cpu().numpy()) _, initial_rank = index.search(target_features.cpu().numpy(), k1) else: # CPU index = index_init_cpu(target_features.size(-1)) index.add(target_features.cpu().numpy()) _, initial_rank = index.search(target_features.cpu().numpy(), k1) nn_k1 = [] nn_k1_half = [] for i in range(N): nn_k1.append(k_reciprocal_neigh(initial_rank, i, k1)) nn_k1_half.append(k_reciprocal_neigh(initial_rank, i, int(np.around(k1/2)))) V = np.zeros((N, N), dtype=mat_type) for i in range(N): k_reciprocal_index = nn_k1[i] k_reciprocal_expansion_index = k_reciprocal_index for candidate in k_reciprocal_index: candidate_k_reciprocal_index = nn_k1_half[candidate] if (len(np.intersect1d(candidate_k_reciprocal_index,k_reciprocal_index)) > 2/3*len(candidate_k_reciprocal_index)): k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index,candidate_k_reciprocal_index) k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) ## element-wise unique dist = 2-2*torch.mm(target_features[i].unsqueeze(0).contiguous(), target_features[k_reciprocal_expansion_index].t()) if use_float16: V[i,k_reciprocal_expansion_index] = F.softmax(-dist, dim=1).view(-1).cpu().numpy().astype(mat_type) else: V[i,k_reciprocal_expansion_index] = F.softmax(-dist, dim=1).view(-1).cpu().numpy() del nn_k1, nn_k1_half if k2 != 1: V_qe = np.zeros_like(V, dtype=mat_type) for i in range(N): V_qe[i,:] = np.mean(V[initial_rank[i,:k2],:], axis=0) V = V_qe del V_qe del initial_rank invIndex = [] for i in range(N): invIndex.append(np.where(V[:,i] != 0)[0]) #len(invIndex)=all_num jaccard_dist = np.zeros((N, N), dtype=mat_type) for i in range(N): temp_min = np.zeros((1, N), dtype=mat_type) # temp_max = np.zeros((1,N), dtype=mat_type) indNonZero = np.where(V[i, :] != 0)[0] indImages = [] indImages = [invIndex[ind] for ind in indNonZero] for j in range(len(indNonZero)): temp_min[0, indImages[j]] = temp_min[0, indImages[j]]+np.minimum(V[i, indNonZero[j]], V[indImages[j], indNonZero[j]]) # temp_max[0,indImages[j]] = temp_max[0,indImages[j]]+np.maximum(V[i,indNonZero[j]],V[indImages[j],indNonZero[j]]) jaccard_dist[i] = 1-temp_min/(2-temp_min) # jaccard_dist[i] = 1-temp_min/(temp_max+1e-6) del invIndex, V pos_bool = (jaccard_dist < 0) jaccard_dist[pos_bool] = 0.0 if print_flag: print("Jaccard distance computing time cost: {}".format(time.time()-end)) return jaccard_dist ================================================ FILE: hhcl/utils/faiss_utils.py ================================================ import os import numpy as np import faiss import torch def swig_ptr_from_FloatTensor(x): assert x.is_contiguous() assert x.dtype == torch.float32 return faiss.cast_integer_to_float_ptr( x.storage().data_ptr() + x.storage_offset() * 4) def swig_ptr_from_LongTensor(x): assert x.is_contiguous() assert x.dtype == torch.int64, 'dtype=%s' % x.dtype return faiss.cast_integer_to_long_ptr( x.storage().data_ptr() + x.storage_offset() * 8) def search_index_pytorch(index, x, k, D=None, I=None): """call the search function of an index with pytorch tensor I/O (CPU and GPU supported)""" assert x.is_contiguous() n, d = x.size() assert d == index.d if D is None: D = torch.empty((n, k), dtype=torch.float32, device=x.device) else: assert D.size() == (n, k) if I is None: I = torch.empty((n, k), dtype=torch.int64, device=x.device) else: assert I.size() == (n, k) torch.cuda.synchronize() xptr = swig_ptr_from_FloatTensor(x) Iptr = swig_ptr_from_LongTensor(I) Dptr = swig_ptr_from_FloatTensor(D) index.search_c(n, xptr, k, Dptr, Iptr) torch.cuda.synchronize() return D, I def search_raw_array_pytorch(res, xb, xq, k, D=None, I=None, metric=faiss.METRIC_L2): assert xb.device == xq.device nq, d = xq.size() if xq.is_contiguous(): xq_row_major = True elif xq.t().is_contiguous(): xq = xq.t() # I initially wrote xq:t(), Lua is still haunting me :-) xq_row_major = False else: raise TypeError('matrix should be row or column-major') xq_ptr = swig_ptr_from_FloatTensor(xq) nb, d2 = xb.size() assert d2 == d if xb.is_contiguous(): xb_row_major = True elif xb.t().is_contiguous(): xb = xb.t() xb_row_major = False else: raise TypeError('matrix should be row or column-major') xb_ptr = swig_ptr_from_FloatTensor(xb) if D is None: D = torch.empty(nq, k, device=xb.device, dtype=torch.float32) else: assert D.shape == (nq, k) assert D.device == xb.device if I is None: I = torch.empty(nq, k, device=xb.device, dtype=torch.int64) else: assert I.shape == (nq, k) assert I.device == xb.device D_ptr = swig_ptr_from_FloatTensor(D) I_ptr = swig_ptr_from_LongTensor(I) faiss.bruteForceKnn(res, metric, xb_ptr, xb_row_major, nb, xq_ptr, xq_row_major, nq, d, k, D_ptr, I_ptr) return D, I def index_init_gpu(ngpus, feat_dim): flat_config = [] for i in range(ngpus): cfg = faiss.GpuIndexFlatConfig() cfg.useFloat16 = False cfg.device = i flat_config.append(cfg) res = [faiss.StandardGpuResources() for i in range(ngpus)] indexes = [faiss.GpuIndexFlatL2(res[i], feat_dim, flat_config[i]) for i in range(ngpus)] index = faiss.IndexShards(feat_dim) for sub_index in indexes: index.add_shard(sub_index) index.reset() return index def index_init_cpu(feat_dim): return faiss.IndexFlatL2(feat_dim) ================================================ FILE: hhcl/utils/logging.py ================================================ from __future__ import absolute_import import os import sys from .osutils import mkdir_if_missing class Logger(object): def __init__(self, fpath=None): self.console = sys.stdout self.file = None if fpath is not None: mkdir_if_missing(os.path.dirname(fpath)) self.file = open(fpath, 'w') def __del__(self): self.close() def __enter__(self): pass def __exit__(self, *args): self.close() def write(self, msg): self.console.write(msg) if self.file is not None: self.file.write(msg) def flush(self): self.console.flush() if self.file is not None: self.file.flush() os.fsync(self.file.fileno()) def close(self): self.console.close() if self.file is not None: self.file.close() ================================================ FILE: hhcl/utils/meters.py ================================================ from __future__ import absolute_import class AverageMeter(object): """Computes and stores the average and current value""" def __init__(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def reset(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def update(self, val, n=1): self.val = val self.sum += val * n self.count += n self.avg = self.sum / self.count ================================================ FILE: hhcl/utils/osutils.py ================================================ from __future__ import absolute_import import os import errno def mkdir_if_missing(dir_path): try: os.makedirs(dir_path) except OSError as e: if e.errno != errno.EEXIST: raise ================================================ FILE: hhcl/utils/rerank.py ================================================ #!/usr/bin/env python2/python3 # -*- coding: utf-8 -*- """ Source: https://github.com/zhunzhong07/person-re-ranking Created on Mon Jun 26 14:46:56 2017 @author: luohao Modified by Houjing Huang, 2017-12-22. - This version accepts distance matrix instead of raw features. - The difference of `/` division between python 2 and 3 is handled. - numpy.float16 is replaced by numpy.float32 for numerical precision. CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017. url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf Matlab version: https://github.com/zhunzhong07/person-re-ranking API q_g_dist: query-gallery distance matrix, numpy array, shape [num_query, num_gallery] q_q_dist: query-query distance matrix, numpy array, shape [num_query, num_query] g_g_dist: gallery-gallery distance matrix, numpy array, shape [num_gallery, num_gallery] k1, k2, lambda_value: parameters, the original paper is (k1=20, k2=6, lambda_value=0.3) Returns: final_dist: re-ranked distance, numpy array, shape [num_query, num_gallery] """ from __future__ import absolute_import from __future__ import print_function from __future__ import division __all__ = ['re_ranking'] import numpy as np def re_ranking(q_g_dist, q_q_dist, g_g_dist, k1=20, k2=6, lambda_value=0.3): # The following naming, e.g. gallery_num, is different from outer scope. # Don't care about it. original_dist = np.concatenate( [np.concatenate([q_q_dist, q_g_dist], axis=1), np.concatenate([q_g_dist.T, g_g_dist], axis=1)], axis=0) original_dist = np.power(original_dist, 2).astype(np.float32) original_dist = np.transpose(1. * original_dist/np.max(original_dist,axis = 0)) V = np.zeros_like(original_dist).astype(np.float32) initial_rank = np.argsort(original_dist).astype(np.int32) query_num = q_g_dist.shape[0] gallery_num = q_g_dist.shape[0] + q_g_dist.shape[1] all_num = gallery_num for i in range(all_num): # k-reciprocal neighbors forward_k_neigh_index = initial_rank[i,:k1+1] backward_k_neigh_index = initial_rank[forward_k_neigh_index,:k1+1] fi = np.where(backward_k_neigh_index==i)[0] k_reciprocal_index = forward_k_neigh_index[fi] k_reciprocal_expansion_index = k_reciprocal_index for j in range(len(k_reciprocal_index)): candidate = k_reciprocal_index[j] candidate_forward_k_neigh_index = initial_rank[candidate,:int(np.around(k1/2.))+1] candidate_backward_k_neigh_index = initial_rank[candidate_forward_k_neigh_index,:int(np.around(k1/2.))+1] fi_candidate = np.where(candidate_backward_k_neigh_index == candidate)[0] candidate_k_reciprocal_index = candidate_forward_k_neigh_index[fi_candidate] if len(np.intersect1d(candidate_k_reciprocal_index,k_reciprocal_index))> 2./3*len(candidate_k_reciprocal_index): k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index,candidate_k_reciprocal_index) k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) weight = np.exp(-original_dist[i,k_reciprocal_expansion_index]) V[i,k_reciprocal_expansion_index] = 1.*weight/np.sum(weight) original_dist = original_dist[:query_num,] if k2 != 1: V_qe = np.zeros_like(V,dtype=np.float32) for i in range(all_num): V_qe[i,:] = np.mean(V[initial_rank[i,:k2],:],axis=0) V = V_qe del V_qe del initial_rank invIndex = [] for i in range(gallery_num): invIndex.append(np.where(V[:,i] != 0)[0]) jaccard_dist = np.zeros_like(original_dist,dtype = np.float32) for i in range(query_num): temp_min = np.zeros(shape=[1,gallery_num],dtype=np.float32) indNonZero = np.where(V[i,:] != 0)[0] indImages = [] indImages = [invIndex[ind] for ind in indNonZero] for j in range(len(indNonZero)): temp_min[0,indImages[j]] = temp_min[0,indImages[j]]+ np.minimum(V[i,indNonZero[j]],V[indImages[j],indNonZero[j]]) jaccard_dist[i] = 1-temp_min/(2.-temp_min) final_dist = jaccard_dist*(1-lambda_value) + original_dist*lambda_value del original_dist del V del jaccard_dist final_dist = final_dist[:query_num,query_num:] return final_dist ================================================ FILE: hhcl/utils/serialization.py ================================================ from __future__ import print_function, absolute_import import json import os.path as osp import shutil import torch from torch.nn import Parameter from .osutils import mkdir_if_missing def read_json(fpath): with open(fpath, 'r') as f: obj = json.load(f) return obj def write_json(obj, fpath): mkdir_if_missing(osp.dirname(fpath)) with open(fpath, 'w') as f: json.dump(obj, f, indent=4, separators=(',', ': ')) def save_checkpoint(state, is_best, fpath='checkpoint.pth.tar'): mkdir_if_missing(osp.dirname(fpath)) torch.save(state, fpath) if is_best: shutil.copy(fpath, osp.join(osp.dirname(fpath), 'model_best.pth.tar')) def load_checkpoint(fpath): if osp.isfile(fpath): # checkpoint = torch.load(fpath) checkpoint = torch.load(fpath, map_location=torch.device('cpu')) print("=> Loaded checkpoint '{}'".format(fpath)) return checkpoint else: raise ValueError("=> No checkpoint found at '{}'".format(fpath)) def copy_state_dict(state_dict, model, strip=None): tgt_state = model.state_dict() copied_names = set() for name, param in state_dict.items(): if strip is not None and name.startswith(strip): name = name[len(strip):] if name not in tgt_state: continue if isinstance(param, Parameter): param = param.data if param.size() != tgt_state[name].size(): print('mismatch:', name, param.size(), tgt_state[name].size()) continue tgt_state[name].copy_(param) copied_names.add(name) missing = set(tgt_state.keys()) - copied_names if len(missing) > 0: print("missing keys in state_dict:", missing) return model ================================================ FILE: requirements.txt ================================================ numpy sklearn Cython h5py pyzmq pillow-simd six scipy matplotlib faiss-gpu==1.6.3 easydict ================================================ FILE: run.sh ================================================ ### resnet50 ### # market1501 CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet50 -d market1501 --iters 200 --eps 0.45 --num-instances 16 --pooling-type avg --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/market1501/resnet50_avg_cmhybrid # dukemtmcreid CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet50 -d dukemtmcreid --iters 200 --eps 0.6 --num-instances 16 --pooling-type avg --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/dukemtmcreid/resnet50_avg_cmhybrid ### resnet_ibn50a + gem pooling ### # market1501 CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet_ibn50a -d market1501 --iters 200 --eps 0.45 --num-instances 16 --pooling-type gem --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/market1501/resnet50_ibn_gem_cmhybrid # dukemtmcreid CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet_ibn50a -d dukemtmcreid --iters 200 --eps 0.6 --num-instances 16 --pooling-type gem --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/dukemtmcreid/resnet50_ibn_gem_cmhybrid # test CUDA_VISIBLE_DEVICES=0 python examples/test.py -d market1501 --data-dir examples/data/market1501 --pooling-type avg --resume examples/logs/market1501/resnet50_avg_cmhybrid/model_best.pth.tar ================================================ FILE: setup.py ================================================ from setuptools import setup, find_packages setup(name='hhcl', version='1.0.0', install_requires=[ 'numpy', 'torch', 'torchvision', 'six', 'h5py', 'Pillow', 'scipy', 'scikit-learn', 'metric-learn', 'faiss_gpu'], packages=find_packages(), keywords=[ 'Unsupervised Learning', 'Contrastive Learning', 'Object Re-identification' ])