mirror of https://github.com/JDAI-CV/fast-reid.git
onnx/trt support
Summary: change model pretrain mode and support onnx/TensorRT exportpull/224/head
liaoxingyu
parent
ee634df290
commit
16655448c2
|
|
@ -4,6 +4,7 @@ FastReID is a research platform that implements state-of-the-art re-identificati
|
|||
|
||||
## What's New
|
||||
|
||||
- [Aug 2020] ONNX/TensorRT converter is supported.
|
||||
- [Jul 2020] Distributed training with multiple GPUs, it trains much faster.
|
||||
- [Jul 2020] `MAX_ITER` in config means `epoch`, it will auto scale to maximum iterations.
|
||||
- Includes more features such as circle loss, abundant visualization methods and evaluation metrics, SoTA results on conventional, cross-domain, partial and vehicle re-id, testing on multi-datasets simultaneously, etc.
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ MODEL:
|
|||
HEADS:
|
||||
NECK_FEAT: "after"
|
||||
POOL_LAYER: "gempool"
|
||||
CLS_LAYER: "circle"
|
||||
CLS_LAYER: "circleSoftmax"
|
||||
SCALE: 64
|
||||
MARGIN: 0.35
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ MODEL:
|
|||
BACKBONE:
|
||||
NAME: "build_resnet_backbone"
|
||||
NORM: "BN"
|
||||
DEPTH: 50
|
||||
DEPTH: "50x"
|
||||
LAST_STRIDE: 1
|
||||
WITH_IBN: False
|
||||
PRETRAIN: True
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ MODEL:
|
|||
BACKBONE:
|
||||
DEPTH: 101
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet101_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("DukeMTMC",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-AGW.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("DukeMTMC",)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ MODEL:
|
|||
BACKBONE:
|
||||
DEPTH: 101
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet101_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("DukeMTMC",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-bagtricks.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("DukeMTMC",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-MGN.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("DukeMTMC",)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ MODEL:
|
|||
BACKBONE:
|
||||
DEPTH: 101
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet101_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("DukeMTMC",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-Strongerbaseline.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("DukeMTMC",)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ MODEL:
|
|||
BACKBONE:
|
||||
DEPTH: 101
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet101_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("MSMT17",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-AGW.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("MSMT17",)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ MODEL:
|
|||
BACKBONE:
|
||||
DEPTH: 101
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet101_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("MSMT17",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-bagtricks.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("MSMT17",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-MGN.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("MSMT17",)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ MODEL:
|
|||
BACKBONE:
|
||||
DEPTH: 101
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet101_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("MSMT17",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-Strongerbaseline.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("MSMT17",)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ MODEL:
|
|||
BACKBONE:
|
||||
DEPTH: 101
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet101_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("Market1501",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-AGW.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("Market1501",)
|
||||
|
|
|
|||
|
|
@ -2,9 +2,8 @@ _BASE_: "../Base-bagtricks.yml"
|
|||
|
||||
MODEL:
|
||||
BACKBONE:
|
||||
DEPTH: 101
|
||||
DEPTH: "101"
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet101_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("Market1501",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-bagtricks.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("Market1501",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-MGN.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("Market1501",)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ MODEL:
|
|||
BACKBONE:
|
||||
DEPTH: 101
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet101_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("Market1501",)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ _BASE_: "../Base-Strongerbaseline.yml"
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/home/liaoxingyu2/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
DATASETS:
|
||||
NAMES: ("Market1501",)
|
||||
|
|
|
|||
|
|
@ -7,9 +7,7 @@ INPUT:
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: '/export2/home/zjk/pretrain_models/resnet50_ibn_a.pth.tar'
|
||||
HEADS:
|
||||
NUM_CLASSES: 30671
|
||||
POOL_LAYER: gempool
|
||||
LOSSES:
|
||||
TRI:
|
||||
|
|
|
|||
|
|
@ -7,10 +7,6 @@ INPUT:
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/export2/home/zjk/pretrain_models/resnet50_ibn_a.pth.tar"
|
||||
|
||||
HEADS:
|
||||
NUM_CLASSES: 575
|
||||
|
||||
SOLVER:
|
||||
OPT: "SGD"
|
||||
|
|
|
|||
|
|
@ -7,9 +7,7 @@ INPUT:
|
|||
MODEL:
|
||||
BACKBONE:
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: '/export2/home/zjk/pretrain_models/resnet50_ibn_a.pth.tar'
|
||||
HEADS:
|
||||
NUM_CLASSES: 13164
|
||||
POOL_LAYER: gempool
|
||||
LOSSES:
|
||||
TRI:
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ cudnn.benchmark = True
|
|||
def setup_cfg(args):
|
||||
# load config from file and command-line arguments
|
||||
cfg = get_cfg()
|
||||
# add_partialreid_config(cfg)
|
||||
cfg.merge_from_file(args.config_file)
|
||||
cfg.merge_from_list(args.opts)
|
||||
cfg.freeze()
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import atexit
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
|
|
|||
|
|
@ -30,9 +30,7 @@ _C.MODEL.FREEZE_LAYERS = ['']
|
|||
_C.MODEL.BACKBONE = CN()
|
||||
|
||||
_C.MODEL.BACKBONE.NAME = "build_resnet_backbone"
|
||||
_C.MODEL.BACKBONE.DEPTH = 50
|
||||
# RegNet volume
|
||||
_C.MODEL.BACKBONE.VOLUME = "800y"
|
||||
_C.MODEL.BACKBONE.DEPTH = "50x"
|
||||
_C.MODEL.BACKBONE.LAST_STRIDE = 1
|
||||
# Normalization method for the convolution layers.
|
||||
_C.MODEL.BACKBONE.NORM = "BN"
|
||||
|
|
@ -137,7 +135,13 @@ _C.INPUT.DO_PAD = True
|
|||
_C.INPUT.PADDING_MODE = 'constant'
|
||||
_C.INPUT.PADDING = 10
|
||||
# Random color jitter
|
||||
_C.INPUT.DO_CJ = False
|
||||
_C.INPUT.CJ = CN()
|
||||
_C.INPUT.CJ.ENABLED = False
|
||||
_C.INPUT.CJ.PROB = 0.8
|
||||
_C.INPUT.CJ.BRIGHTNESS = 0.15
|
||||
_C.INPUT.CJ.CONTRAST = 0.15
|
||||
_C.INPUT.CJ.SATURATION = 0.1
|
||||
_C.INPUT.CJ.HUE = 0.1
|
||||
# Auto augmentation
|
||||
_C.INPUT.DO_AUTOAUG = False
|
||||
# Augmix augmentation
|
||||
|
|
|
|||
|
|
@ -1,124 +1,124 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: Jinkai Zheng
|
||||
@contact: 1315673509@qq.com
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
import random
|
||||
|
||||
from .bases import ImageDataset
|
||||
from ..datasets import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VehicleID(ImageDataset):
|
||||
"""VehicleID.
|
||||
|
||||
Reference:
|
||||
Liu et al. Deep relative distance learning: Tell the difference between similar vehicles. CVPR 2016.
|
||||
|
||||
URL: `<https://pkuml.org/resources/pku-vehicleid.html>`_
|
||||
|
||||
Train dataset statistics:
|
||||
- identities: 13164.
|
||||
- images: 113346.
|
||||
"""
|
||||
dataset_dir = "vehicleid"
|
||||
dataset_name = "vehicleid"
|
||||
|
||||
def __init__(self, root='datasets', test_list='', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
|
||||
self.image_dir = osp.join(self.dataset_dir, 'image')
|
||||
self.train_list = osp.join(self.dataset_dir, 'train_test_split/train_list.txt')
|
||||
if test_list:
|
||||
self.test_list = test_list
|
||||
else:
|
||||
self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_13164.txt')
|
||||
|
||||
required_files = [
|
||||
self.dataset_dir,
|
||||
self.image_dir,
|
||||
self.train_list,
|
||||
self.test_list,
|
||||
]
|
||||
self.check_before_run(required_files)
|
||||
|
||||
train = self.process_dir(self.train_list, is_train=True)
|
||||
query, gallery = self.process_dir(self.test_list, is_train=False)
|
||||
|
||||
super(VehicleID, self).__init__(train, query, gallery, **kwargs)
|
||||
|
||||
def process_dir(self, list_file, is_train=True):
|
||||
img_list_lines = open(list_file, 'r').readlines()
|
||||
|
||||
dataset = []
|
||||
for idx, line in enumerate(img_list_lines):
|
||||
line = line.strip()
|
||||
vid = int(line.split(' ')[1])
|
||||
imgid = line.split(' ')[0]
|
||||
img_path = osp.join(self.image_dir, imgid + '.jpg')
|
||||
if is_train:
|
||||
vid = self.dataset_name + "_" + str(vid)
|
||||
dataset.append((img_path, vid, int(imgid)))
|
||||
|
||||
if is_train: return dataset
|
||||
else:
|
||||
random.shuffle(dataset)
|
||||
vid_container = set()
|
||||
query = []
|
||||
gallery = []
|
||||
for sample in dataset:
|
||||
if sample[1] not in vid_container:
|
||||
vid_container.add(sample[1])
|
||||
gallery.append(sample)
|
||||
else:
|
||||
query.append(sample)
|
||||
|
||||
return query, gallery
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class SmallVehicleID(VehicleID):
|
||||
"""VehicleID.
|
||||
Small test dataset statistics:
|
||||
- identities: 800.
|
||||
- images: 6493.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_800.txt')
|
||||
|
||||
super(SmallVehicleID, self).__init__(root, self.test_list, **kwargs)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class MediumVehicleID(VehicleID):
|
||||
"""VehicleID.
|
||||
Medium test dataset statistics:
|
||||
- identities: 1600.
|
||||
- images: 13377.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_1600.txt')
|
||||
|
||||
super(MediumVehicleID, self).__init__(root, self.test_list, **kwargs)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class LargeVehicleID(VehicleID):
|
||||
"""VehicleID.
|
||||
Large test dataset statistics:
|
||||
- identities: 2400.
|
||||
- images: 19777.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_2400.txt')
|
||||
|
||||
super(LargeVehicleID, self).__init__(root, self.test_list, **kwargs)
|
||||
# encoding: utf-8
|
||||
"""
|
||||
@author: Jinkai Zheng
|
||||
@contact: 1315673509@qq.com
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
import random
|
||||
|
||||
from .bases import ImageDataset
|
||||
from ..datasets import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VehicleID(ImageDataset):
|
||||
"""VehicleID.
|
||||
|
||||
Reference:
|
||||
Liu et al. Deep relative distance learning: Tell the difference between similar vehicles. CVPR 2016.
|
||||
|
||||
URL: `<https://pkuml.org/resources/pku-vehicleid.html>`_
|
||||
|
||||
Train dataset statistics:
|
||||
- identities: 13164.
|
||||
- images: 113346.
|
||||
"""
|
||||
dataset_dir = "vehicleid"
|
||||
dataset_name = "vehicleid"
|
||||
|
||||
def __init__(self, root='datasets', test_list='', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
|
||||
self.image_dir = osp.join(self.dataset_dir, 'image')
|
||||
self.train_list = osp.join(self.dataset_dir, 'train_test_split/train_list.txt')
|
||||
if test_list:
|
||||
self.test_list = test_list
|
||||
else:
|
||||
self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_13164.txt')
|
||||
|
||||
required_files = [
|
||||
self.dataset_dir,
|
||||
self.image_dir,
|
||||
self.train_list,
|
||||
self.test_list,
|
||||
]
|
||||
self.check_before_run(required_files)
|
||||
|
||||
train = self.process_dir(self.train_list, is_train=True)
|
||||
query, gallery = self.process_dir(self.test_list, is_train=False)
|
||||
|
||||
super(VehicleID, self).__init__(train, query, gallery, **kwargs)
|
||||
|
||||
def process_dir(self, list_file, is_train=True):
|
||||
img_list_lines = open(list_file, 'r').readlines()
|
||||
|
||||
dataset = []
|
||||
for idx, line in enumerate(img_list_lines):
|
||||
line = line.strip()
|
||||
vid = int(line.split(' ')[1])
|
||||
imgid = line.split(' ')[0]
|
||||
img_path = osp.join(self.image_dir, imgid + '.jpg')
|
||||
if is_train:
|
||||
vid = self.dataset_name + "_" + str(vid)
|
||||
dataset.append((img_path, vid, int(imgid)))
|
||||
|
||||
if is_train: return dataset
|
||||
else:
|
||||
random.shuffle(dataset)
|
||||
vid_container = set()
|
||||
query = []
|
||||
gallery = []
|
||||
for sample in dataset:
|
||||
if sample[1] not in vid_container:
|
||||
vid_container.add(sample[1])
|
||||
gallery.append(sample)
|
||||
else:
|
||||
query.append(sample)
|
||||
|
||||
return query, gallery
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class SmallVehicleID(VehicleID):
|
||||
"""VehicleID.
|
||||
Small test dataset statistics:
|
||||
- identities: 800.
|
||||
- images: 6493.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_800.txt')
|
||||
|
||||
super(SmallVehicleID, self).__init__(root, self.test_list, **kwargs)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class MediumVehicleID(VehicleID):
|
||||
"""VehicleID.
|
||||
Medium test dataset statistics:
|
||||
- identities: 1600.
|
||||
- images: 13377.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_1600.txt')
|
||||
|
||||
super(MediumVehicleID, self).__init__(root, self.test_list, **kwargs)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class LargeVehicleID(VehicleID):
|
||||
"""VehicleID.
|
||||
Large test dataset statistics:
|
||||
- identities: 2400.
|
||||
- images: 19777.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_2400.txt')
|
||||
|
||||
super(LargeVehicleID, self).__init__(root, self.test_list, **kwargs)
|
||||
|
|
|
|||
|
|
@ -1,67 +1,67 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: Jinkai Zheng
|
||||
@contact: 1315673509@qq.com
|
||||
"""
|
||||
|
||||
import glob
|
||||
import os.path as osp
|
||||
import re
|
||||
|
||||
from .bases import ImageDataset
|
||||
from ..datasets import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VeRi(ImageDataset):
|
||||
"""VeRi.
|
||||
|
||||
Reference:
|
||||
Liu et al. A Deep Learning based Approach for Progressive Vehicle Re-Identification. ECCV 2016.
|
||||
|
||||
URL: `<https://vehiclereid.github.io/VeRi/>`_
|
||||
|
||||
Dataset statistics:
|
||||
- identities: 775.
|
||||
- images: 37778 (train) + 1678 (query) + 11579 (gallery).
|
||||
"""
|
||||
dataset_dir = "veri"
|
||||
dataset_name = "veri"
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
|
||||
self.train_dir = osp.join(self.dataset_dir, 'image_train')
|
||||
self.query_dir = osp.join(self.dataset_dir, 'image_query')
|
||||
self.gallery_dir = osp.join(self.dataset_dir, 'image_test')
|
||||
|
||||
required_files = [
|
||||
self.dataset_dir,
|
||||
self.train_dir,
|
||||
self.query_dir,
|
||||
self.gallery_dir,
|
||||
]
|
||||
self.check_before_run(required_files)
|
||||
|
||||
train = self.process_dir(self.train_dir)
|
||||
query = self.process_dir(self.query_dir, is_train=False)
|
||||
gallery = self.process_dir(self.gallery_dir, is_train=False)
|
||||
|
||||
super(VeRi, self).__init__(train, query, gallery, **kwargs)
|
||||
|
||||
def process_dir(self, dir_path, is_train=True):
|
||||
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
|
||||
pattern = re.compile(r'([\d]+)_c(\d\d\d)')
|
||||
|
||||
data = []
|
||||
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 <= pid <= 776
|
||||
assert 1 <= camid <= 20
|
||||
camid -= 1 # index starts from 0
|
||||
if is_train:
|
||||
pid = self.dataset_name + "_" + str(pid)
|
||||
data.append((img_path, pid, camid))
|
||||
|
||||
return data
|
||||
# encoding: utf-8
|
||||
"""
|
||||
@author: Jinkai Zheng
|
||||
@contact: 1315673509@qq.com
|
||||
"""
|
||||
|
||||
import glob
|
||||
import os.path as osp
|
||||
import re
|
||||
|
||||
from .bases import ImageDataset
|
||||
from ..datasets import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VeRi(ImageDataset):
|
||||
"""VeRi.
|
||||
|
||||
Reference:
|
||||
Liu et al. A Deep Learning based Approach for Progressive Vehicle Re-Identification. ECCV 2016.
|
||||
|
||||
URL: `<https://vehiclereid.github.io/VeRi/>`_
|
||||
|
||||
Dataset statistics:
|
||||
- identities: 775.
|
||||
- images: 37778 (train) + 1678 (query) + 11579 (gallery).
|
||||
"""
|
||||
dataset_dir = "veri"
|
||||
dataset_name = "veri"
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
|
||||
self.train_dir = osp.join(self.dataset_dir, 'image_train')
|
||||
self.query_dir = osp.join(self.dataset_dir, 'image_query')
|
||||
self.gallery_dir = osp.join(self.dataset_dir, 'image_test')
|
||||
|
||||
required_files = [
|
||||
self.dataset_dir,
|
||||
self.train_dir,
|
||||
self.query_dir,
|
||||
self.gallery_dir,
|
||||
]
|
||||
self.check_before_run(required_files)
|
||||
|
||||
train = self.process_dir(self.train_dir)
|
||||
query = self.process_dir(self.query_dir, is_train=False)
|
||||
gallery = self.process_dir(self.gallery_dir, is_train=False)
|
||||
|
||||
super(VeRi, self).__init__(train, query, gallery, **kwargs)
|
||||
|
||||
def process_dir(self, dir_path, is_train=True):
|
||||
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
|
||||
pattern = re.compile(r'([\d]+)_c(\d\d\d)')
|
||||
|
||||
data = []
|
||||
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 <= pid <= 776
|
||||
assert 1 <= camid <= 20
|
||||
camid -= 1 # index starts from 0
|
||||
if is_train:
|
||||
pid = self.dataset_name + "_" + str(pid)
|
||||
data.append((img_path, pid, camid))
|
||||
|
||||
return data
|
||||
|
|
|
|||
|
|
@ -1,138 +1,138 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: Jinkai Zheng
|
||||
@contact: 1315673509@qq.com
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
|
||||
from .bases import ImageDataset
|
||||
from ..datasets import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VeRiWild(ImageDataset):
|
||||
"""VeRi-Wild.
|
||||
|
||||
Reference:
|
||||
Lou et al. A Large-Scale Dataset for Vehicle Re-Identification in the Wild. CVPR 2019.
|
||||
|
||||
URL: `<https://github.com/PKU-IMRE/VERI-Wild>`_
|
||||
|
||||
Train dataset statistics:
|
||||
- identities: 30671.
|
||||
- images: 277797.
|
||||
"""
|
||||
dataset_dir = "VERI-Wild"
|
||||
dataset_name = "veriwild"
|
||||
|
||||
def __init__(self, root='datasets', query_list='', gallery_list='', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
|
||||
self.image_dir = osp.join(self.dataset_dir, 'images')
|
||||
self.train_list = osp.join(self.dataset_dir, 'train_test_split/train_list.txt')
|
||||
self.vehicle_info = osp.join(self.dataset_dir, 'train_test_split/vehicle_info.txt')
|
||||
if query_list and gallery_list:
|
||||
self.query_list = query_list
|
||||
self.gallery_list = gallery_list
|
||||
else:
|
||||
self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_10000_query.txt')
|
||||
self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_10000.txt')
|
||||
|
||||
required_files = [
|
||||
self.image_dir,
|
||||
self.train_list,
|
||||
self.query_list,
|
||||
self.gallery_list,
|
||||
self.vehicle_info,
|
||||
]
|
||||
self.check_before_run(required_files)
|
||||
|
||||
self.imgid2vid, self.imgid2camid, self.imgid2imgpath = self.process_vehicle(self.vehicle_info)
|
||||
|
||||
train = self.process_dir(self.train_list)
|
||||
query = self.process_dir(self.query_list, is_train=False)
|
||||
gallery = self.process_dir(self.gallery_list, is_train=False)
|
||||
|
||||
super(VeRiWild, self).__init__(train, query, gallery, **kwargs)
|
||||
|
||||
def process_dir(self, img_list, is_train=True):
|
||||
img_list_lines = open(img_list, 'r').readlines()
|
||||
|
||||
dataset = []
|
||||
for idx, line in enumerate(img_list_lines):
|
||||
line = line.strip()
|
||||
vid = int(line.split('/')[0])
|
||||
imgid = line.split('/')[1]
|
||||
if is_train:
|
||||
vid = self.dataset_name + "_" + str(vid)
|
||||
dataset.append((self.imgid2imgpath[imgid], vid, int(self.imgid2camid[imgid])))
|
||||
|
||||
assert len(dataset) == len(img_list_lines)
|
||||
return dataset
|
||||
|
||||
def process_vehicle(self, vehicle_info):
|
||||
imgid2vid = {}
|
||||
imgid2camid = {}
|
||||
imgid2imgpath = {}
|
||||
vehicle_info_lines = open(vehicle_info, 'r').readlines()
|
||||
|
||||
for idx, line in enumerate(vehicle_info_lines[1:]):
|
||||
vid = line.strip().split('/')[0]
|
||||
imgid = line.strip().split(';')[0].split('/')[1]
|
||||
camid = line.strip().split(';')[1]
|
||||
img_path = osp.join(self.image_dir, vid, imgid + '.jpg')
|
||||
imgid2vid[imgid] = vid
|
||||
imgid2camid[imgid] = camid
|
||||
imgid2imgpath[imgid] = img_path
|
||||
|
||||
assert len(imgid2vid) == len(vehicle_info_lines) - 1
|
||||
return imgid2vid, imgid2camid, imgid2imgpath
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class SmallVeRiWild(VeRiWild):
|
||||
"""VeRi-Wild.
|
||||
Small test dataset statistics:
|
||||
- identities: 3000.
|
||||
- images: 41861.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_3000_query.txt')
|
||||
self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_3000.txt')
|
||||
|
||||
super(SmallVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class MediumVeRiWild(VeRiWild):
|
||||
"""VeRi-Wild.
|
||||
Medium test dataset statistics:
|
||||
- identities: 5000.
|
||||
- images: 69389.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_5000_query.txt')
|
||||
self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_5000.txt')
|
||||
|
||||
super(MediumVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class LargeVeRiWild(VeRiWild):
|
||||
"""VeRi-Wild.
|
||||
Large test dataset statistics:
|
||||
- identities: 10000.
|
||||
- images: 138517.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_10000_query.txt')
|
||||
self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_10000.txt')
|
||||
|
||||
super(LargeVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
|
||||
# encoding: utf-8
|
||||
"""
|
||||
@author: Jinkai Zheng
|
||||
@contact: 1315673509@qq.com
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
|
||||
from .bases import ImageDataset
|
||||
from ..datasets import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VeRiWild(ImageDataset):
|
||||
"""VeRi-Wild.
|
||||
|
||||
Reference:
|
||||
Lou et al. A Large-Scale Dataset for Vehicle Re-Identification in the Wild. CVPR 2019.
|
||||
|
||||
URL: `<https://github.com/PKU-IMRE/VERI-Wild>`_
|
||||
|
||||
Train dataset statistics:
|
||||
- identities: 30671.
|
||||
- images: 277797.
|
||||
"""
|
||||
dataset_dir = "VERI-Wild"
|
||||
dataset_name = "veriwild"
|
||||
|
||||
def __init__(self, root='datasets', query_list='', gallery_list='', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
|
||||
self.image_dir = osp.join(self.dataset_dir, 'images')
|
||||
self.train_list = osp.join(self.dataset_dir, 'train_test_split/train_list.txt')
|
||||
self.vehicle_info = osp.join(self.dataset_dir, 'train_test_split/vehicle_info.txt')
|
||||
if query_list and gallery_list:
|
||||
self.query_list = query_list
|
||||
self.gallery_list = gallery_list
|
||||
else:
|
||||
self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_10000_query.txt')
|
||||
self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_10000.txt')
|
||||
|
||||
required_files = [
|
||||
self.image_dir,
|
||||
self.train_list,
|
||||
self.query_list,
|
||||
self.gallery_list,
|
||||
self.vehicle_info,
|
||||
]
|
||||
self.check_before_run(required_files)
|
||||
|
||||
self.imgid2vid, self.imgid2camid, self.imgid2imgpath = self.process_vehicle(self.vehicle_info)
|
||||
|
||||
train = self.process_dir(self.train_list)
|
||||
query = self.process_dir(self.query_list, is_train=False)
|
||||
gallery = self.process_dir(self.gallery_list, is_train=False)
|
||||
|
||||
super(VeRiWild, self).__init__(train, query, gallery, **kwargs)
|
||||
|
||||
def process_dir(self, img_list, is_train=True):
|
||||
img_list_lines = open(img_list, 'r').readlines()
|
||||
|
||||
dataset = []
|
||||
for idx, line in enumerate(img_list_lines):
|
||||
line = line.strip()
|
||||
vid = int(line.split('/')[0])
|
||||
imgid = line.split('/')[1]
|
||||
if is_train:
|
||||
vid = self.dataset_name + "_" + str(vid)
|
||||
dataset.append((self.imgid2imgpath[imgid], vid, int(self.imgid2camid[imgid])))
|
||||
|
||||
assert len(dataset) == len(img_list_lines)
|
||||
return dataset
|
||||
|
||||
def process_vehicle(self, vehicle_info):
|
||||
imgid2vid = {}
|
||||
imgid2camid = {}
|
||||
imgid2imgpath = {}
|
||||
vehicle_info_lines = open(vehicle_info, 'r').readlines()
|
||||
|
||||
for idx, line in enumerate(vehicle_info_lines[1:]):
|
||||
vid = line.strip().split('/')[0]
|
||||
imgid = line.strip().split(';')[0].split('/')[1]
|
||||
camid = line.strip().split(';')[1]
|
||||
img_path = osp.join(self.image_dir, vid, imgid + '.jpg')
|
||||
imgid2vid[imgid] = vid
|
||||
imgid2camid[imgid] = camid
|
||||
imgid2imgpath[imgid] = img_path
|
||||
|
||||
assert len(imgid2vid) == len(vehicle_info_lines) - 1
|
||||
return imgid2vid, imgid2camid, imgid2imgpath
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class SmallVeRiWild(VeRiWild):
|
||||
"""VeRi-Wild.
|
||||
Small test dataset statistics:
|
||||
- identities: 3000.
|
||||
- images: 41861.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_3000_query.txt')
|
||||
self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_3000.txt')
|
||||
|
||||
super(SmallVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class MediumVeRiWild(VeRiWild):
|
||||
"""VeRi-Wild.
|
||||
Medium test dataset statistics:
|
||||
- identities: 5000.
|
||||
- images: 69389.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_5000_query.txt')
|
||||
self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_5000.txt')
|
||||
|
||||
super(MediumVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class LargeVeRiWild(VeRiWild):
|
||||
"""VeRi-Wild.
|
||||
Large test dataset statistics:
|
||||
- identities: 10000.
|
||||
- images: 138517.
|
||||
"""
|
||||
|
||||
def __init__(self, root='datasets', **kwargs):
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_10000_query.txt')
|
||||
self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_10000.txt')
|
||||
|
||||
super(LargeVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
|
||||
|
|
|
|||
|
|
@ -33,7 +33,12 @@ def build_transforms(cfg, is_train=True):
|
|||
padding_mode = cfg.INPUT.PADDING_MODE
|
||||
|
||||
# color jitter
|
||||
do_cj = cfg.INPUT.DO_CJ
|
||||
do_cj = cfg.INPUT.CJ.ENABLED
|
||||
cj_prob = cfg.INPUT.CJ.PROB
|
||||
cj_brightness = cfg.INPUT.CJ.BRIGHTNESS
|
||||
cj_contrast = cfg.INPUT.CJ.CONTRAST
|
||||
cj_saturation = cfg.INPUT.CJ.SATURATION
|
||||
cj_hue = cfg.INPUT.CJ.HUE
|
||||
|
||||
# random erasing
|
||||
do_rea = cfg.INPUT.REA.ENABLED
|
||||
|
|
@ -52,7 +57,7 @@ def build_transforms(cfg, is_train=True):
|
|||
res.extend([T.Pad(padding, padding_mode=padding_mode),
|
||||
T.RandomCrop(size_train)])
|
||||
if do_cj:
|
||||
res.append(T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0))
|
||||
T.RandomApply([T.ColorJitter(cj_brightness, cj_contrast, cj_saturation, cj_hue)], p=cj_prob)
|
||||
if do_augmix:
|
||||
res.append(AugMix())
|
||||
if do_rea:
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
__all__ = ['ToTensor', 'RandomErasing', 'RandomPatch', 'AugMix', ]
|
||||
__all__ = ['ToTensor', 'RandomErasing', 'RandomPatch', 'AugMix',]
|
||||
|
||||
import math
|
||||
import random
|
||||
|
|
@ -202,110 +202,3 @@ class AugMix(object):
|
|||
|
||||
mixed = (1 - m) * image + m * mix
|
||||
return mixed
|
||||
|
||||
# class ColorJitter(object):
|
||||
# """docstring for do_color"""
|
||||
#
|
||||
# def __init__(self, probability=0.5):
|
||||
# self.probability = probability
|
||||
#
|
||||
# def do_brightness_shift(self, image, alpha=0.125):
|
||||
# image = image.astype(np.float32)
|
||||
# image = image + alpha * 255
|
||||
# image = np.clip(image, 0, 255).astype(np.uint8)
|
||||
# return image
|
||||
#
|
||||
# def do_brightness_multiply(self, image, alpha=1):
|
||||
# image = image.astype(np.float32)
|
||||
# image = alpha * image
|
||||
# image = np.clip(image, 0, 255).astype(np.uint8)
|
||||
# return image
|
||||
#
|
||||
# def do_contrast(self, image, alpha=1.0):
|
||||
# image = image.astype(np.float32)
|
||||
# gray = image * np.array([[[0.114, 0.587, 0.299]]]) # rgb to gray (YCbCr)
|
||||
# gray = (3.0 * (1.0 - alpha) / gray.size) * np.sum(gray)
|
||||
# image = alpha * image + gray
|
||||
# image = np.clip(image, 0, 255).astype(np.uint8)
|
||||
# return image
|
||||
#
|
||||
# # https://www.pyimagesearch.com/2015/10/05/opencv-gamma-correction/
|
||||
# def do_gamma(self, image, gamma=1.0):
|
||||
# table = np.array([((i / 255.0) ** (1.0 / gamma)) * 255
|
||||
# for i in np.arange(0, 256)]).astype("uint8")
|
||||
#
|
||||
# return cv2.LUT(image, table) # apply gamma correction using the lookup table
|
||||
#
|
||||
# def do_clahe(self, image, clip=2, grid=16):
|
||||
# grid = int(grid)
|
||||
#
|
||||
# lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
|
||||
# gray, a, b = cv2.split(lab)
|
||||
# gray = cv2.createCLAHE(clipLimit=clip, tileGridSize=(grid, grid)).apply(gray)
|
||||
# lab = cv2.merge((gray, a, b))
|
||||
# image = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
|
||||
#
|
||||
# return image
|
||||
#
|
||||
# def __call__(self, image):
|
||||
# if random.uniform(0, 1) > self.probability:
|
||||
# return image
|
||||
#
|
||||
# image = np.asarray(image, dtype=np.uint8).copy()
|
||||
# index = random.randint(0, 4)
|
||||
# if index == 0:
|
||||
# image = self.do_brightness_shift(image, 0.1)
|
||||
# elif index == 1:
|
||||
# image = self.do_gamma(image, 1)
|
||||
# elif index == 2:
|
||||
# image = self.do_clahe(image)
|
||||
# elif index == 3:
|
||||
# image = self.do_brightness_multiply(image)
|
||||
# elif index == 4:
|
||||
# image = self.do_contrast(image)
|
||||
# return image
|
||||
|
||||
|
||||
# class random_shift(object):
|
||||
# """docstring for do_color"""
|
||||
#
|
||||
# def __init__(self, probability=0.5):
|
||||
# self.probability = probability
|
||||
#
|
||||
# def __call__(self, image):
|
||||
# if random.uniform(0, 1) > self.probability:
|
||||
# return image
|
||||
#
|
||||
# width, height, d = image.shape
|
||||
# zero_image = np.zeros_like(image)
|
||||
# w = random.randint(0, 20) - 10
|
||||
# h = random.randint(0, 30) - 15
|
||||
# zero_image[max(0, w): min(w + width, width), max(h, 0): min(h + height, height)] = \
|
||||
# image[max(0, -w): min(-w + width, width), max(-h, 0): min(-h + height, height)]
|
||||
# image = zero_image.copy()
|
||||
# return image
|
||||
#
|
||||
#
|
||||
# class random_scale(object):
|
||||
# """docstring for do_color"""
|
||||
#
|
||||
# def __init__(self, probability=0.5):
|
||||
# self.probability = probability
|
||||
#
|
||||
# def __call__(self, image):
|
||||
# if random.uniform(0, 1) > self.probability:
|
||||
# return image
|
||||
#
|
||||
# scale = random.random() * 0.1 + 0.9
|
||||
# assert 0.9 <= scale <= 1
|
||||
# width, height, d = image.shape
|
||||
# zero_image = np.zeros_like(image)
|
||||
# new_width = round(width * scale)
|
||||
# new_height = round(height * scale)
|
||||
# image = cv2.resize(image, (new_height, new_width))
|
||||
# start_w = random.randint(0, width - new_width)
|
||||
# start_h = random.randint(0, height - new_height)
|
||||
# zero_image[start_w: start_w + new_width,
|
||||
# start_h:start_h + new_height] = image
|
||||
# image = zero_image.copy()
|
||||
# return image
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
# based on:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
# based on
|
||||
|
|
|
|||
|
|
@ -10,7 +10,6 @@ from collections import OrderedDict
|
|||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from tabulate import tabulate
|
||||
|
||||
from .evaluator import DatasetEvaluator
|
||||
from .query_expansion import aqe
|
||||
|
|
@ -101,6 +100,6 @@ class ReidEvaluator(DatasetEvaluator):
|
|||
tprs = evaluate_roc(dist, query_pids, gallery_pids, query_camids, gallery_camids)
|
||||
fprs = [1e-4, 1e-3, 1e-2]
|
||||
for i in range(len(fprs)):
|
||||
self._results["TPR@FPR={}".format(fprs[i])] = tprs[i]
|
||||
self._results["TPR@FPR={:.0e}".format(fprs[i])] = tprs[i]
|
||||
|
||||
return copy.deepcopy(self._results)
|
||||
|
|
|
|||
|
|
@ -6,9 +6,10 @@
|
|||
|
||||
from .activation import *
|
||||
from .arc_softmax import ArcSoftmax
|
||||
from .circle_softmax import CircleSoftmax
|
||||
from .am_softmax import AMSoftmax
|
||||
from .batch_drop import BatchDrop
|
||||
from .batch_norm import *
|
||||
from .circle_softmax import CircleSoftmax
|
||||
from .context_block import ContextBlock
|
||||
from .frn import FRN, TLU
|
||||
from .non_local import Non_local
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import math
|
||||
|
|
|
|||
|
|
@ -0,0 +1,43 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn import Parameter
|
||||
|
||||
|
||||
class AMSoftmax(nn.Module):
|
||||
r"""Implement of large margin cosine distance:
|
||||
Args:
|
||||
in_feat: size of each input sample
|
||||
num_classes: size of each output sample
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, in_feat, num_classes):
|
||||
super().__init__()
|
||||
self.in_features = in_feat
|
||||
self._num_classes = num_classes
|
||||
self._s = cfg.MODEL.HEADS.SCALE
|
||||
self._m = cfg.MODEL.HEADS.MARGIN
|
||||
self.weight = Parameter(torch.Tensor(num_classes, in_feat))
|
||||
nn.init.xavier_uniform_(self.weight)
|
||||
|
||||
def forward(self, features, targets):
|
||||
# --------------------------- cos(theta) & phi(theta) ---------------------------
|
||||
cosine = F.linear(F.normalize(features), F.normalize(self.weight))
|
||||
phi = cosine - self._m
|
||||
# --------------------------- convert label to one-hot ---------------------------
|
||||
targets = F.one_hot(targets, num_classes=self._num_classes)
|
||||
output = (targets * phi) + ((1.0 - targets) * cosine)
|
||||
output *= self._s
|
||||
|
||||
return output
|
||||
|
||||
def extra_repr(self):
|
||||
return 'in_features={}, num_classes={}, scale={}, margin={}'.format(
|
||||
self.in_feat, self._num_classes, self._s, self._m
|
||||
)
|
||||
|
|
@ -26,6 +26,7 @@ class ArcSoftmax(nn.Module):
|
|||
self.mm = math.sin(math.pi - self._m) * self._m
|
||||
|
||||
self.weight = Parameter(torch.Tensor(num_classes, in_feat))
|
||||
nn.init.xavier_uniform_(self.weight)
|
||||
self.register_buffer('t', torch.zeros(1))
|
||||
|
||||
def forward(self, features, targets):
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@
|
|||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
|
@ -19,11 +21,12 @@ class CircleSoftmax(nn.Module):
|
|||
self._m = cfg.MODEL.HEADS.MARGIN
|
||||
|
||||
self.weight = Parameter(torch.Tensor(num_classes, in_feat))
|
||||
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
|
||||
def forward(self, features, targets):
|
||||
sim_mat = F.linear(F.normalize(features), F.normalize(self.weight))
|
||||
alpha_p = F.relu(-sim_mat.detach() + 1 + self._m)
|
||||
alpha_n = F.relu(sim_mat.detach() + self._m)
|
||||
alpha_p = torch.clamp_min(-sim_mat.detach() + 1 + self._m, min=0.)
|
||||
alpha_n = torch.clamp_min(sim_mat.detach() + self._m, min=0.)
|
||||
delta_p = 1 - self._m
|
||||
delta_n = self._m
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,4 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
|
|
|||
|
|
@ -10,4 +10,4 @@ from .resnet import build_resnet_backbone
|
|||
from .osnet import build_osnet_backbone
|
||||
from .resnest import build_resnest_backbone
|
||||
from .resnext import build_resnext_backbone
|
||||
from .regnet import build_regnet_backbone
|
||||
from .regnet import build_regnet_backbone
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
# based on:
|
||||
|
|
@ -9,7 +9,8 @@
|
|||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from fastreid.layers import get_norm
|
||||
from .build import BACKBONE_REGISTRY
|
||||
|
||||
model_urls = {
|
||||
|
|
@ -37,6 +38,8 @@ class ConvLayer(nn.Module):
|
|||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
bn_norm,
|
||||
num_splits,
|
||||
stride=1,
|
||||
padding=0,
|
||||
groups=1,
|
||||
|
|
@ -55,7 +58,7 @@ class ConvLayer(nn.Module):
|
|||
if IN:
|
||||
self.bn = nn.InstanceNorm2d(out_channels, affine=True)
|
||||
else:
|
||||
self.bn = nn.BatchNorm2d(out_channels)
|
||||
self.bn = get_norm(bn_norm, out_channels, num_splits)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
|
|
@ -68,7 +71,7 @@ class ConvLayer(nn.Module):
|
|||
class Conv1x1(nn.Module):
|
||||
"""1x1 convolution + bn + relu."""
|
||||
|
||||
def __init__(self, in_channels, out_channels, stride=1, groups=1):
|
||||
def __init__(self, in_channels, out_channels, bn_norm, num_splits, stride=1, groups=1):
|
||||
super(Conv1x1, self).__init__()
|
||||
self.conv = nn.Conv2d(
|
||||
in_channels,
|
||||
|
|
@ -79,7 +82,7 @@ class Conv1x1(nn.Module):
|
|||
bias=False,
|
||||
groups=groups
|
||||
)
|
||||
self.bn = nn.BatchNorm2d(out_channels)
|
||||
self.bn = get_norm(bn_norm, out_channels, num_splits)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
|
|
@ -92,12 +95,12 @@ class Conv1x1(nn.Module):
|
|||
class Conv1x1Linear(nn.Module):
|
||||
"""1x1 convolution + bn (w/o non-linearity)."""
|
||||
|
||||
def __init__(self, in_channels, out_channels, stride=1):
|
||||
def __init__(self, in_channels, out_channels, bn_norm, num_splits, stride=1):
|
||||
super(Conv1x1Linear, self).__init__()
|
||||
self.conv = nn.Conv2d(
|
||||
in_channels, out_channels, 1, stride=stride, padding=0, bias=False
|
||||
)
|
||||
self.bn = nn.BatchNorm2d(out_channels)
|
||||
self.bn = get_norm(bn_norm, out_channels, num_splits)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
|
|
@ -108,7 +111,7 @@ class Conv1x1Linear(nn.Module):
|
|||
class Conv3x3(nn.Module):
|
||||
"""3x3 convolution + bn + relu."""
|
||||
|
||||
def __init__(self, in_channels, out_channels, stride=1, groups=1):
|
||||
def __init__(self, in_channels, out_channels, bn_norm, num_splits, stride=1, groups=1):
|
||||
super(Conv3x3, self).__init__()
|
||||
self.conv = nn.Conv2d(
|
||||
in_channels,
|
||||
|
|
@ -119,7 +122,7 @@ class Conv3x3(nn.Module):
|
|||
bias=False,
|
||||
groups=groups
|
||||
)
|
||||
self.bn = nn.BatchNorm2d(out_channels)
|
||||
self.bn = get_norm(bn_norm, out_channels, num_splits)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
|
|
@ -134,7 +137,7 @@ class LightConv3x3(nn.Module):
|
|||
1x1 (linear) + dw 3x3 (nonlinear).
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels):
|
||||
def __init__(self, in_channels, out_channels, bn_norm, num_splits):
|
||||
super(LightConv3x3, self).__init__()
|
||||
self.conv1 = nn.Conv2d(
|
||||
in_channels, out_channels, 1, stride=1, padding=0, bias=False
|
||||
|
|
@ -148,7 +151,7 @@ class LightConv3x3(nn.Module):
|
|||
bias=False,
|
||||
groups=out_channels
|
||||
)
|
||||
self.bn = nn.BatchNorm2d(out_channels)
|
||||
self.bn = get_norm(bn_norm, out_channels, num_splits)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
|
|
@ -197,9 +200,12 @@ class ChannelGate(nn.Module):
|
|||
bias=True,
|
||||
padding=0
|
||||
)
|
||||
if gate_activation == 'sigmoid': self.gate_activation = nn.Sigmoid()
|
||||
elif gate_activation == 'relu': self.gate_activation = nn.ReLU(inplace=True)
|
||||
elif gate_activation == 'linear': self.gate_activation = nn.Identity()
|
||||
if gate_activation == 'sigmoid':
|
||||
self.gate_activation = nn.Sigmoid()
|
||||
elif gate_activation == 'relu':
|
||||
self.gate_activation = nn.ReLU(inplace=True)
|
||||
elif gate_activation == 'linear':
|
||||
self.gate_activation = nn.Identity()
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Unknown gate activation: {}".format(gate_activation)
|
||||
|
|
@ -224,34 +230,36 @@ class OSBlock(nn.Module):
|
|||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
bn_norm,
|
||||
num_splits,
|
||||
IN=False,
|
||||
bottleneck_reduction=4,
|
||||
**kwargs
|
||||
):
|
||||
super(OSBlock, self).__init__()
|
||||
mid_channels = out_channels // bottleneck_reduction
|
||||
self.conv1 = Conv1x1(in_channels, mid_channels)
|
||||
self.conv2a = LightConv3x3(mid_channels, mid_channels)
|
||||
self.conv1 = Conv1x1(in_channels, mid_channels, bn_norm, num_splits)
|
||||
self.conv2a = LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits)
|
||||
self.conv2b = nn.Sequential(
|
||||
LightConv3x3(mid_channels, mid_channels),
|
||||
LightConv3x3(mid_channels, mid_channels),
|
||||
LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits),
|
||||
LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits),
|
||||
)
|
||||
self.conv2c = nn.Sequential(
|
||||
LightConv3x3(mid_channels, mid_channels),
|
||||
LightConv3x3(mid_channels, mid_channels),
|
||||
LightConv3x3(mid_channels, mid_channels),
|
||||
LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits),
|
||||
LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits),
|
||||
LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits),
|
||||
)
|
||||
self.conv2d = nn.Sequential(
|
||||
LightConv3x3(mid_channels, mid_channels),
|
||||
LightConv3x3(mid_channels, mid_channels),
|
||||
LightConv3x3(mid_channels, mid_channels),
|
||||
LightConv3x3(mid_channels, mid_channels),
|
||||
LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits),
|
||||
LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits),
|
||||
LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits),
|
||||
LightConv3x3(mid_channels, mid_channels, bn_norm, num_splits),
|
||||
)
|
||||
self.gate = ChannelGate(mid_channels)
|
||||
self.conv3 = Conv1x1Linear(mid_channels, out_channels)
|
||||
self.conv3 = Conv1x1Linear(mid_channels, out_channels, bn_norm, num_splits)
|
||||
self.downsample = None
|
||||
if in_channels != out_channels:
|
||||
self.downsample = Conv1x1Linear(in_channels, out_channels)
|
||||
self.downsample = Conv1x1Linear(in_channels, out_channels, bn_norm, num_splits)
|
||||
self.IN = None
|
||||
if IN: self.IN = nn.InstanceNorm2d(out_channels, affine=True)
|
||||
self.relu = nn.ReLU(True)
|
||||
|
|
@ -290,6 +298,8 @@ class OSNet(nn.Module):
|
|||
blocks,
|
||||
layers,
|
||||
channels,
|
||||
bn_norm,
|
||||
num_splits,
|
||||
IN=False,
|
||||
**kwargs
|
||||
):
|
||||
|
|
@ -299,13 +309,15 @@ class OSNet(nn.Module):
|
|||
assert num_blocks == len(channels) - 1
|
||||
|
||||
# convolutional backbone
|
||||
self.conv1 = ConvLayer(3, channels[0], 7, stride=2, padding=3, IN=IN)
|
||||
self.conv1 = ConvLayer(3, channels[0], 7, bn_norm, num_splits, stride=2, padding=3, IN=IN)
|
||||
self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
|
||||
self.conv2 = self._make_layer(
|
||||
blocks[0],
|
||||
layers[0],
|
||||
channels[0],
|
||||
channels[1],
|
||||
bn_norm,
|
||||
num_splits,
|
||||
reduce_spatial_size=True,
|
||||
IN=IN
|
||||
)
|
||||
|
|
@ -314,6 +326,8 @@ class OSNet(nn.Module):
|
|||
layers[1],
|
||||
channels[1],
|
||||
channels[2],
|
||||
bn_norm,
|
||||
num_splits,
|
||||
reduce_spatial_size=True
|
||||
)
|
||||
self.conv4 = self._make_layer(
|
||||
|
|
@ -321,9 +335,11 @@ class OSNet(nn.Module):
|
|||
layers[2],
|
||||
channels[2],
|
||||
channels[3],
|
||||
bn_norm,
|
||||
num_splits,
|
||||
reduce_spatial_size=False
|
||||
)
|
||||
self.conv5 = Conv1x1(channels[3], channels[3])
|
||||
self.conv5 = Conv1x1(channels[3], channels[3], bn_norm, num_splits)
|
||||
|
||||
self._init_params()
|
||||
|
||||
|
|
@ -333,19 +349,21 @@ class OSNet(nn.Module):
|
|||
layer,
|
||||
in_channels,
|
||||
out_channels,
|
||||
bn_norm,
|
||||
num_splits,
|
||||
reduce_spatial_size,
|
||||
IN=False
|
||||
):
|
||||
layers = []
|
||||
|
||||
layers.append(block(in_channels, out_channels, IN=IN))
|
||||
layers.append(block(in_channels, out_channels, bn_norm, num_splits, IN=IN))
|
||||
for i in range(1, layer):
|
||||
layers.append(block(out_channels, out_channels, IN=IN))
|
||||
layers.append(block(out_channels, out_channels, bn_norm, num_splits, IN=IN))
|
||||
|
||||
if reduce_spatial_size:
|
||||
layers.append(
|
||||
nn.Sequential(
|
||||
Conv1x1(out_channels, out_channels),
|
||||
Conv1x1(out_channels, out_channels, bn_norm, num_splits),
|
||||
nn.AvgPool2d(2, stride=2),
|
||||
)
|
||||
)
|
||||
|
|
@ -477,11 +495,19 @@ def build_osnet_backbone(cfg):
|
|||
# fmt: off
|
||||
pretrain = cfg.MODEL.BACKBONE.PRETRAIN
|
||||
with_ibn = cfg.MODEL.BACKBONE.WITH_IBN
|
||||
bn_norm = cfg.MODEL.BACKBONE.NORM
|
||||
num_splits = cfg.MODEL.BACKBONE.NORM_SPLIT
|
||||
depth = cfg.MODEL.BACKBONE.DEPTH
|
||||
|
||||
num_blocks_per_stage = [2, 2, 2]
|
||||
num_channels_per_stage = [64, 256, 384, 512]
|
||||
model = OSNet([OSBlock, OSBlock, OSBlock], num_blocks_per_stage, num_channels_per_stage, with_ibn)
|
||||
pretrain_key = 'osnet_ibn_x1_0' if with_ibn else 'osnet_x1_0'
|
||||
num_channels_per_stage = {"x1_0": [64, 256, 384, 512], "x0_75": [48, 192, 288, 384], "x0_5": [32, 128, 192, 256],
|
||||
"x0_25": [16, 64, 96, 128]}[depth]
|
||||
model = OSNet([OSBlock, OSBlock, OSBlock], num_blocks_per_stage, num_channels_per_stage,
|
||||
bn_norm, num_splits, IN=with_ibn)
|
||||
|
||||
if pretrain:
|
||||
if with_ibn: pretrain_key = "osnet_ibn_" + depth
|
||||
else: pretrain_key = "osnet_" + depth
|
||||
|
||||
init_pretrained_weights(model, pretrain_key)
|
||||
return model
|
||||
|
|
|
|||
|
|
@ -10,7 +10,18 @@ from ..build import BACKBONE_REGISTRY
|
|||
from .config import regnet_cfg
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
model_urls = {
|
||||
'800x': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/160906036/RegNetX-800MF_dds_8gpu.pyth',
|
||||
'800y': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/160906567/RegNetY-800MF_dds_8gpu.pyth',
|
||||
'1600x': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/160990626/RegNetX-1.6GF_dds_8gpu.pyth',
|
||||
'1600y': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/160906681/RegNetY-1.6GF_dds_8gpu.pyth',
|
||||
'3200x': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/160906139/RegNetX-3.2GF_dds_8gpu.pyth',
|
||||
'3200y': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/160906834/RegNetY-3.2GF_dds_8gpu.pyth',
|
||||
'4000x': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/160906383/RegNetX-4.0GF_dds_8gpu.pyth',
|
||||
'4000y': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/160906838/RegNetY-4.0GF_dds_8gpu.pyth',
|
||||
'6400x': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/161116590/RegNetX-6.4GF_dds_8gpu.pyth',
|
||||
'6400y': 'https://dl.fbaipublicfiles.com/pycls/dds_baselines/160907112/RegNetY-6.4GF_dds_8gpu.pyth',
|
||||
}
|
||||
|
||||
def init_weights(m):
|
||||
"""Performs ResNet-style weight initialization."""
|
||||
|
|
@ -464,14 +475,60 @@ class RegNet(AnyNet):
|
|||
super(RegNet, self).__init__(**kwargs)
|
||||
|
||||
|
||||
def init_pretrained_weights(key):
|
||||
"""Initializes model with pretrained weights.
|
||||
|
||||
Layers that don't match with pretrained layers in name or size are kept unchanged.
|
||||
"""
|
||||
import os
|
||||
import errno
|
||||
import gdown
|
||||
|
||||
def _get_torch_home():
|
||||
ENV_TORCH_HOME = 'TORCH_HOME'
|
||||
ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'
|
||||
DEFAULT_CACHE_DIR = '~/.cache'
|
||||
torch_home = os.path.expanduser(
|
||||
os.getenv(
|
||||
ENV_TORCH_HOME,
|
||||
os.path.join(
|
||||
os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'
|
||||
)
|
||||
)
|
||||
)
|
||||
return torch_home
|
||||
|
||||
torch_home = _get_torch_home()
|
||||
model_dir = os.path.join(torch_home, 'checkpoints')
|
||||
try:
|
||||
os.makedirs(model_dir)
|
||||
except OSError as e:
|
||||
if e.errno == errno.EEXIST:
|
||||
# Directory already exists, ignore.
|
||||
pass
|
||||
else:
|
||||
# Unexpected OSError, re-raise.
|
||||
raise
|
||||
|
||||
filename = model_urls[key].split('/')[-1]
|
||||
|
||||
cached_file = os.path.join(model_dir, filename)
|
||||
|
||||
if not os.path.exists(cached_file):
|
||||
gdown.download(model_urls[key], cached_file, quiet=False)
|
||||
|
||||
logger.info(f"Loading pretrained model from {cached_file}")
|
||||
state_dict = torch.load(cached_file)['model_state']
|
||||
|
||||
return state_dict
|
||||
|
||||
@BACKBONE_REGISTRY.register()
|
||||
def build_regnet_backbone(cfg):
|
||||
# fmt: off
|
||||
pretrain = cfg.MODEL.BACKBONE.PRETRAIN
|
||||
pretrain_path = cfg.MODEL.BACKBONE.PRETRAIN_PATH
|
||||
last_stride = cfg.MODEL.BACKBONE.LAST_STRIDE
|
||||
bn_norm = cfg.MODEL.BACKBONE.NORM
|
||||
volume = cfg.MODEL.BACKBONE.VOLUME
|
||||
depth = cfg.MODEL.BACKBONE.DEPTH
|
||||
|
||||
cfg_files = {
|
||||
'800x': 'fastreid/modeling/backbones/regnet/regnetx/RegNetX-800MF_dds_8gpu.yaml',
|
||||
|
|
@ -480,21 +537,18 @@ def build_regnet_backbone(cfg):
|
|||
'1600y': 'fastreid/modeling/backbones/regnet/regnety/RegNetY-1.6GF_dds_8gpu.yaml',
|
||||
'3200x': 'fastreid/modeling/backbones/regnet/regnetx/RegNetX-3.2GF_dds_8gpu.yaml',
|
||||
'3200y': 'fastreid/modeling/backbones/regnet/regnety/RegNetY-3.2GF_dds_8gpu.yaml',
|
||||
'4000x': 'fastreid/modeling/backbones/regnet/regnety/RegNetX-4.0GF_dds_8gpu.yaml',
|
||||
'4000y': 'fastreid/modeling/backbones/regnet/regnety/RegNetY-4.0GF_dds_8gpu.yaml',
|
||||
'6400x': 'fastreid/modeling/backbones/regnet/regnetx/RegNetX-6.4GF_dds_8gpu.yaml',
|
||||
'6400y': 'fastreid/modeling/backbones/regnet/regnety/RegNetY-6.4GF_dds_8gpu.yaml',
|
||||
}[volume]
|
||||
}[depth]
|
||||
|
||||
regnet_cfg.merge_from_file(cfg_files)
|
||||
model = RegNet(last_stride, bn_norm)
|
||||
|
||||
if pretrain:
|
||||
try:
|
||||
state_dict = torch.load(pretrain_path, map_location=torch.device('cpu'))['model_state']
|
||||
except FileNotFoundError as e:
|
||||
logger.info(f'{pretrain_path} is not found! Please check this path.')
|
||||
raise e
|
||||
|
||||
logger.info(f"Loading pretrained model from {pretrain_path}")
|
||||
key = depth
|
||||
state_dict = init_pretrained_weights(key)
|
||||
|
||||
incompatible = model.load_state_dict(state_dict, strict=False)
|
||||
if incompatible.missing_keys:
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ import math
|
|||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils import model_zoo
|
||||
|
||||
from fastreid.layers import (
|
||||
IBN,
|
||||
|
|
@ -22,11 +21,15 @@ from .build import BACKBONE_REGISTRY
|
|||
|
||||
logger = logging.getLogger(__name__)
|
||||
model_urls = {
|
||||
18: 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||
34: 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||
50: 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||
101: 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||
152: 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
||||
'18x': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||
'34x': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||
'50x': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||
'101x': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||
'ibn_18x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet18_ibn_a-2f571257.pth',
|
||||
'ibn_34x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet34_ibn_a-94bc1577.pth',
|
||||
'ibn_50x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet50_ibn_a-d9d0bb7b.pth',
|
||||
'ibn_101x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet101_ibn_a-59ea0ac6.pth',
|
||||
'se_ibn_101x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/se_resnet101_ibn_a-fabed4e2.pth',
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -37,7 +40,10 @@ class BasicBlock(nn.Module):
|
|||
stride=1, downsample=None, reduction=16):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
|
||||
self.bn1 = get_norm(bn_norm, planes, num_splits)
|
||||
if with_ibn:
|
||||
self.bn1 = IBN(planes, bn_norm, num_splits)
|
||||
else:
|
||||
self.bn1 = get_norm(bn_norm, planes, num_splits)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn2 = get_norm(bn_norm, planes, num_splits)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
|
@ -146,8 +152,6 @@ class ResNet(nn.Module):
|
|||
)
|
||||
|
||||
layers = []
|
||||
if planes == 512:
|
||||
with_ibn = False
|
||||
layers.append(block(self.inplanes, planes, bn_norm, num_splits, with_ibn, with_se, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
|
|
@ -227,6 +231,54 @@ class ResNet(nn.Module):
|
|||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
|
||||
def init_pretrained_weights(key):
|
||||
"""Initializes model with pretrained weights.
|
||||
|
||||
Layers that don't match with pretrained layers in name or size are kept unchanged.
|
||||
"""
|
||||
import os
|
||||
import errno
|
||||
import gdown
|
||||
|
||||
def _get_torch_home():
|
||||
ENV_TORCH_HOME = 'TORCH_HOME'
|
||||
ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'
|
||||
DEFAULT_CACHE_DIR = '~/.cache'
|
||||
torch_home = os.path.expanduser(
|
||||
os.getenv(
|
||||
ENV_TORCH_HOME,
|
||||
os.path.join(
|
||||
os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'
|
||||
)
|
||||
)
|
||||
)
|
||||
return torch_home
|
||||
|
||||
torch_home = _get_torch_home()
|
||||
model_dir = os.path.join(torch_home, 'checkpoints')
|
||||
try:
|
||||
os.makedirs(model_dir)
|
||||
except OSError as e:
|
||||
if e.errno == errno.EEXIST:
|
||||
# Directory already exists, ignore.
|
||||
pass
|
||||
else:
|
||||
# Unexpected OSError, re-raise.
|
||||
raise
|
||||
|
||||
filename = model_urls[key].split('/')[-1]
|
||||
|
||||
cached_file = os.path.join(model_dir, filename)
|
||||
|
||||
if not os.path.exists(cached_file):
|
||||
gdown.download(model_urls[key], cached_file, quiet=False)
|
||||
|
||||
logger.info(f"Loading pretrained model from {cached_file}")
|
||||
state_dict = torch.load(cached_file)
|
||||
|
||||
return state_dict
|
||||
|
||||
|
||||
@BACKBONE_REGISTRY.register()
|
||||
def build_resnet_backbone(cfg):
|
||||
"""
|
||||
|
|
@ -246,13 +298,15 @@ def build_resnet_backbone(cfg):
|
|||
with_nl = cfg.MODEL.BACKBONE.WITH_NL
|
||||
depth = cfg.MODEL.BACKBONE.DEPTH
|
||||
|
||||
num_blocks_per_stage = {34: [3, 4, 6, 3], 50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3], }[depth]
|
||||
nl_layers_per_stage = {34: [0, 2, 3, 0], 50: [0, 2, 3, 0], 101: [0, 2, 9, 0]}[depth]
|
||||
block = {34: BasicBlock, 50: Bottleneck, 101: Bottleneck}[depth]
|
||||
num_blocks_per_stage = {'18x': [2, 2, 2, 2], '34x': [3, 4, 6, 3], '50x': [3, 4, 6, 3],
|
||||
'101x': [3, 4, 23, 3],}[depth]
|
||||
nl_layers_per_stage = {'18x': [0, 0, 0, 0], '34x': [0, 0, 0, 0], '50x': [0, 2, 3, 0], '101x': [0, 2, 9, 0]}[depth]
|
||||
block = {'18x': BasicBlock, '34x': BasicBlock, '50x': Bottleneck, '101x': Bottleneck}[depth]
|
||||
model = ResNet(last_stride, bn_norm, num_splits, with_ibn, with_se, with_nl, block,
|
||||
num_blocks_per_stage, nl_layers_per_stage)
|
||||
if pretrain:
|
||||
if not with_ibn:
|
||||
# Load pretrain path if specifically
|
||||
if pretrain_path:
|
||||
try:
|
||||
state_dict = torch.load(pretrain_path, map_location=torch.device('cpu'))['model']
|
||||
# Remove module.encoder in name
|
||||
|
|
@ -263,20 +317,19 @@ def build_resnet_backbone(cfg):
|
|||
new_state_dict[new_k] = state_dict[k]
|
||||
state_dict = new_state_dict
|
||||
logger.info(f"Loading pretrained model from {pretrain_path}")
|
||||
except FileNotFoundError or KeyError:
|
||||
# original resnet
|
||||
state_dict = model_zoo.load_url(model_urls[depth])
|
||||
logger.info("Loading pretrained model from torchvision")
|
||||
except FileNotFoundError as e:
|
||||
logger.info(f'{pretrain_path} is not found! Please check this path.')
|
||||
raise e
|
||||
except KeyError as e:
|
||||
logger.info("State dict keys error! Please check the state dict.")
|
||||
raise e
|
||||
else:
|
||||
state_dict = torch.load(pretrain_path, map_location=torch.device('cpu'))['state_dict'] # ibn-net
|
||||
# Remove module in name
|
||||
new_state_dict = {}
|
||||
for k in state_dict:
|
||||
new_k = '.'.join(k.split('.')[1:])
|
||||
if new_k in model.state_dict() and (model.state_dict()[new_k].shape == state_dict[k].shape):
|
||||
new_state_dict[new_k] = state_dict[k]
|
||||
state_dict = new_state_dict
|
||||
logger.info(f"Loading pretrained model from {pretrain_path}")
|
||||
key = depth
|
||||
if with_ibn: key = 'ibn_' + key
|
||||
if with_se: key = 'se_' + key
|
||||
|
||||
state_dict = init_pretrained_weights(key)
|
||||
|
||||
incompatible = model.load_state_dict(state_dict, strict=False)
|
||||
if incompatible.missing_keys:
|
||||
logger.info(
|
||||
|
|
@ -286,4 +339,5 @@ def build_resnet_backbone(cfg):
|
|||
logger.info(
|
||||
get_unexpected_parameters_message(incompatible.unexpected_keys)
|
||||
)
|
||||
|
||||
return model
|
||||
|
|
|
|||
|
|
@ -1,21 +1,27 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
# based on:
|
||||
# https://github.com/XingangPan/IBN-Net/blob/master/models/imagenet/resnext_ibn_a.py
|
||||
|
||||
import math
|
||||
import logging
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn import init
|
||||
import math
|
||||
|
||||
import torch
|
||||
from ...layers import IBN
|
||||
import torch.nn as nn
|
||||
|
||||
from fastreid.layers import IBN, get_norm
|
||||
from fastreid.utils.checkpoint import get_missing_parameters_message, get_unexpected_parameters_message
|
||||
from .build import BACKBONE_REGISTRY
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
model_urls = {
|
||||
'ibn_101x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnext101_ibn_a-6ace051d.pth',
|
||||
}
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
"""
|
||||
|
|
@ -23,7 +29,8 @@ class Bottleneck(nn.Module):
|
|||
"""
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, with_ibn, baseWidth, cardinality, stride=1, downsample=None):
|
||||
def __init__(self, inplanes, planes, bn_norm, num_splits, with_ibn, baseWidth, cardinality, stride=1,
|
||||
downsample=None):
|
||||
""" Constructor
|
||||
Args:
|
||||
inplanes: input channel dimensionality
|
||||
|
|
@ -38,13 +45,13 @@ class Bottleneck(nn.Module):
|
|||
C = cardinality
|
||||
self.conv1 = nn.Conv2d(inplanes, D * C, kernel_size=1, stride=1, padding=0, bias=False)
|
||||
if with_ibn:
|
||||
self.bn1 = IBN(D * C)
|
||||
self.bn1 = IBN(D * C, bn_norm, num_splits)
|
||||
else:
|
||||
self.bn1 = nn.BatchNorm2d(D * C)
|
||||
self.bn1 = get_norm(bn_norm, D * C, num_splits)
|
||||
self.conv2 = nn.Conv2d(D * C, D * C, kernel_size=3, stride=stride, padding=1, groups=C, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(D * C)
|
||||
self.bn2 = get_norm(bn_norm, D * C, num_splits)
|
||||
self.conv3 = nn.Conv2d(D * C, planes * 4, kernel_size=1, stride=1, padding=0, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * 4)
|
||||
self.bn3 = get_norm(bn_norm, planes * 4, num_splits)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
self.downsample = downsample
|
||||
|
|
@ -78,7 +85,7 @@ class ResNeXt(nn.Module):
|
|||
https://arxiv.org/pdf/1611.05431.pdf
|
||||
"""
|
||||
|
||||
def __init__(self, last_stride, with_ibn, block, layers, baseWidth=4, cardinality=32):
|
||||
def __init__(self, last_stride, bn_norm, num_splits, with_ibn, block, layers, baseWidth=4, cardinality=32):
|
||||
""" Constructor
|
||||
Args:
|
||||
baseWidth: baseWidth for ResNeXt.
|
||||
|
|
@ -94,17 +101,17 @@ class ResNeXt(nn.Module):
|
|||
self.output_size = 64
|
||||
|
||||
self.conv1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.bn1 = get_norm(bn_norm, 64, num_splits)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0], with_ibn=with_ibn)
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2, with_ibn=with_ibn)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2, with_ibn=with_ibn)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=last_stride, with_ibn=with_ibn)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0], 1, bn_norm, num_splits, with_ibn=with_ibn)
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], 2, bn_norm, num_splits, with_ibn=with_ibn)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], 2, bn_norm, num_splits, with_ibn=with_ibn)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], last_stride, bn_norm, num_splits, with_ibn=with_ibn)
|
||||
|
||||
self.random_init()
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1, with_ibn=False):
|
||||
def _make_layer(self, block, planes, blocks, stride=1, bn_norm='BN', num_splits=1, with_ibn=False):
|
||||
""" Stack n bottleneck modules where n is inferred from the depth of the network.
|
||||
Args:
|
||||
block: block type used to construct ResNext
|
||||
|
|
@ -118,16 +125,18 @@ class ResNeXt(nn.Module):
|
|||
downsample = nn.Sequential(
|
||||
nn.Conv2d(self.inplanes, planes * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(planes * block.expansion),
|
||||
get_norm(bn_norm, planes * block.expansion, num_splits),
|
||||
)
|
||||
|
||||
layers = []
|
||||
if planes == 512:
|
||||
with_ibn = False
|
||||
layers.append(block(self.inplanes, planes, with_ibn, self.baseWidth, self.cardinality, stride, downsample))
|
||||
layers.append(block(self.inplanes, planes, bn_norm, num_splits, with_ibn,
|
||||
self.baseWidth, self.cardinality, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, with_ibn, self.baseWidth, self.cardinality, 1, None))
|
||||
layers.append(
|
||||
block(self.inplanes, planes, bn_norm, num_splits, with_ibn, self.baseWidth, self.cardinality, 1, None))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
|
|
@ -157,6 +166,53 @@ class ResNeXt(nn.Module):
|
|||
m.bias.data.zero_()
|
||||
|
||||
|
||||
def init_pretrained_weights(key):
|
||||
"""Initializes model with pretrained weights.
|
||||
|
||||
Layers that don't match with pretrained layers in name or size are kept unchanged.
|
||||
"""
|
||||
import os
|
||||
import errno
|
||||
import gdown
|
||||
|
||||
def _get_torch_home():
|
||||
ENV_TORCH_HOME = 'TORCH_HOME'
|
||||
ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'
|
||||
DEFAULT_CACHE_DIR = '~/.cache'
|
||||
torch_home = os.path.expanduser(
|
||||
os.getenv(
|
||||
ENV_TORCH_HOME,
|
||||
os.path.join(
|
||||
os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'
|
||||
)
|
||||
)
|
||||
)
|
||||
return torch_home
|
||||
|
||||
torch_home = _get_torch_home()
|
||||
model_dir = os.path.join(torch_home, 'checkpoints')
|
||||
try:
|
||||
os.makedirs(model_dir)
|
||||
except OSError as e:
|
||||
if e.errno == errno.EEXIST:
|
||||
# Directory already exists, ignore.
|
||||
pass
|
||||
else:
|
||||
# Unexpected OSError, re-raise.
|
||||
raise
|
||||
|
||||
filename = model_urls[key].split('/')[-1]
|
||||
|
||||
cached_file = os.path.join(model_dir, filename)
|
||||
|
||||
if not os.path.exists(cached_file):
|
||||
gdown.download(model_urls[key], cached_file, quiet=False)
|
||||
|
||||
logger.info(f"Loading pretrained model from {cached_file}")
|
||||
state_dict = torch.load(cached_file)
|
||||
|
||||
return state_dict
|
||||
|
||||
@BACKBONE_REGISTRY.register()
|
||||
def build_resnext_backbone(cfg):
|
||||
"""
|
||||
|
|
@ -169,30 +225,48 @@ def build_resnext_backbone(cfg):
|
|||
pretrain = cfg.MODEL.BACKBONE.PRETRAIN
|
||||
pretrain_path = cfg.MODEL.BACKBONE.PRETRAIN_PATH
|
||||
last_stride = cfg.MODEL.BACKBONE.LAST_STRIDE
|
||||
bn_norm = cfg.MODEL.BACKBONE.NORM
|
||||
num_splits = cfg.MODEL.BACKBONE.NORM_SPLIT
|
||||
with_ibn = cfg.MODEL.BACKBONE.WITH_IBN
|
||||
with_se = cfg.MODEL.BACKBONE.WITH_SE
|
||||
with_nl = cfg.MODEL.BACKBONE.WITH_NL
|
||||
depth = cfg.MODEL.BACKBONE.DEPTH
|
||||
|
||||
num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3], }[depth]
|
||||
nl_layers_per_stage = {50: [0, 2, 3, 0], 101: [0, 2, 3, 0]}[depth]
|
||||
model = ResNeXt(last_stride, with_ibn, Bottleneck, num_blocks_per_stage)
|
||||
num_blocks_per_stage = {'50x': [3, 4, 6, 3], '101x': [3, 4, 23, 3], '152x': [3, 8, 36, 3], }[depth]
|
||||
nl_layers_per_stage = {'50x': [0, 2, 3, 0], '101x': [0, 2, 3, 0]}[depth]
|
||||
model = ResNeXt(last_stride, bn_norm, num_splits, with_ibn, Bottleneck, num_blocks_per_stage)
|
||||
|
||||
if pretrain:
|
||||
# if not with_ibn:
|
||||
# original resnet
|
||||
# state_dict = model_zoo.load_url(model_urls[depth])
|
||||
# else:
|
||||
# ibn resnet
|
||||
state_dict = torch.load(pretrain_path)['state_dict']
|
||||
# remove module in name
|
||||
new_state_dict = {}
|
||||
for k in state_dict:
|
||||
new_k = '.'.join(k.split('.')[1:])
|
||||
if new_k in model.state_dict() and (model.state_dict()[new_k].shape == state_dict[k].shape):
|
||||
new_state_dict[new_k] = state_dict[k]
|
||||
state_dict = new_state_dict
|
||||
res = model.load_state_dict(state_dict, strict=False)
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info('missing keys is {}'.format(res.missing_keys))
|
||||
logger.info('unexpected keys is {}'.format(res.unexpected_keys))
|
||||
if pretrain_path:
|
||||
try:
|
||||
state_dict = torch.load(pretrain_path, map_location=torch.device('cpu'))['model']
|
||||
# Remove module.encoder in name
|
||||
new_state_dict = {}
|
||||
for k in state_dict:
|
||||
new_k = '.'.join(k.split('.')[2:])
|
||||
if new_k in model.state_dict() and (model.state_dict()[new_k].shape == state_dict[k].shape):
|
||||
new_state_dict[new_k] = state_dict[k]
|
||||
state_dict = new_state_dict
|
||||
logger.info(f"Loading pretrained model from {pretrain_path}")
|
||||
except FileNotFoundError as e:
|
||||
logger.info(f'{pretrain_path} is not found! Please check this path.')
|
||||
raise e
|
||||
except KeyError as e:
|
||||
logger.info("State dict keys error! Please check the state dict.")
|
||||
raise e
|
||||
else:
|
||||
key = depth
|
||||
if with_ibn: key = 'ibn_' + key
|
||||
|
||||
state_dict = init_pretrained_weights(key)
|
||||
|
||||
incompatible = model.load_state_dict(state_dict, strict=False)
|
||||
if incompatible.missing_keys:
|
||||
logger.info(
|
||||
get_missing_parameters_message(incompatible.missing_keys)
|
||||
)
|
||||
if incompatible.unexpected_keys:
|
||||
logger.info(
|
||||
get_unexpected_parameters_message(incompatible.unexpected_keys)
|
||||
)
|
||||
|
||||
return model
|
||||
|
|
|
|||
|
|
@ -24,9 +24,10 @@ class BNneckHead(nn.Module):
|
|||
if cls_type == 'linear': self.classifier = nn.Linear(in_feat, num_classes, bias=False)
|
||||
elif cls_type == 'arcSoftmax': self.classifier = ArcSoftmax(cfg, in_feat, num_classes)
|
||||
elif cls_type == 'circleSoftmax': self.classifier = CircleSoftmax(cfg, in_feat, num_classes)
|
||||
elif cls_type == 'amSoftmax': self.classifier = AMSoftmax(cfg, in_feat, num_classes)
|
||||
else:
|
||||
raise KeyError(f"{cls_type} is invalid, please choose from "
|
||||
f"'linear', 'arcSoftmax' and 'circleSoftmax'.")
|
||||
f"'linear', 'arcSoftmax', 'amSoftmax' and 'circleSoftmax'.")
|
||||
|
||||
self.classifier.apply(weights_init_classifier)
|
||||
|
||||
|
|
|
|||
|
|
@ -20,9 +20,10 @@ class LinearHead(nn.Module):
|
|||
if cls_type == 'linear': self.classifier = nn.Linear(in_feat, num_classes, bias=False)
|
||||
elif cls_type == 'arcSoftmax': self.classifier = ArcSoftmax(cfg, in_feat, num_classes)
|
||||
elif cls_type == 'circleSoftmax': self.classifier = CircleSoftmax(cfg, in_feat, num_classes)
|
||||
elif cls_type == 'amSoftmax': self.classifier = AMSoftmax(cfg, in_feat, num_classes)
|
||||
else:
|
||||
raise KeyError(f"{cls_type} is invalid, please choose from "
|
||||
f"'linear', 'arcSoftmax' and 'circleSoftmax'.")
|
||||
f"'linear', 'arcSoftmax', 'amSoftmax' and 'circleSoftmax'.")
|
||||
|
||||
self.classifier.apply(weights_init_classifier)
|
||||
|
||||
|
|
|
|||
|
|
@ -32,12 +32,13 @@ class ReductionHead(nn.Module):
|
|||
|
||||
# identity classification layer
|
||||
cls_type = cfg.MODEL.HEADS.CLS_LAYER
|
||||
if cls_type == 'linear': self.classifier = nn.Linear(in_feat, num_classes, bias=False)
|
||||
elif cls_type == 'arcSoftmax': self.classifier = ArcSoftmax(cfg, in_feat, num_classes)
|
||||
elif cls_type == 'circleSoftmax': self.classifier = CircleSoftmax(cfg, in_feat, num_classes)
|
||||
if cls_type == 'linear': self.classifier = nn.Linear(reduction_dim, num_classes, bias=False)
|
||||
elif cls_type == 'arcSoftmax': self.classifier = ArcSoftmax(cfg, reduction_dim, num_classes)
|
||||
elif cls_type == 'circleSoftmax': self.classifier = CircleSoftmax(cfg, reduction_dim, num_classes)
|
||||
elif cls_type == 'amSoftmax': self.classifier = AMSoftmax(cfg, reduction_dim, num_classes)
|
||||
else:
|
||||
raise KeyError(f"{cls_type} is invalid, please choose from "
|
||||
f"'linear', 'arcSoftmax' and 'circleSoftmax'.")
|
||||
f"'linear', 'arcSoftmax', 'amSoftmax' and 'circleSoftmax'.")
|
||||
|
||||
self.classifier.apply(weights_init_classifier)
|
||||
|
||||
|
|
|
|||
|
|
@ -6,4 +6,5 @@
|
|||
|
||||
from .cross_entroy_loss import CrossEntropyLoss
|
||||
from .focal_loss import FocalLoss
|
||||
from .metric_loss import TripletLoss, CircleLoss
|
||||
from .triplet_loss import TripletLoss
|
||||
from .circle_loss import CircleLoss
|
||||
|
|
|
|||
|
|
@ -0,0 +1,61 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from fastreid.utils import comm
|
||||
from .utils import concat_all_gather
|
||||
|
||||
|
||||
class CircleLoss(object):
|
||||
def __init__(self, cfg):
|
||||
self._scale = cfg.MODEL.LOSSES.CIRCLE.SCALE
|
||||
|
||||
self._m = cfg.MODEL.LOSSES.CIRCLE.MARGIN
|
||||
self._s = cfg.MODEL.LOSSES.CIRCLE.ALPHA
|
||||
|
||||
def __call__(self, embedding, targets):
|
||||
embedding = nn.functional.normalize(embedding, dim=1)
|
||||
|
||||
if comm.get_world_size() > 1:
|
||||
all_embedding = concat_all_gather(embedding)
|
||||
all_targets = concat_all_gather(targets)
|
||||
else:
|
||||
all_embedding = embedding
|
||||
all_targets = targets
|
||||
|
||||
dist_mat = torch.matmul(embedding, all_embedding.t())
|
||||
|
||||
N, M = dist_mat.size()
|
||||
is_pos = targets.view(N, 1).expand(N, M).eq(all_targets.view(M, 1).expand(M, N).t()).float()
|
||||
|
||||
# Compute the mask which ignores the relevance score of the query to itself
|
||||
if M > N:
|
||||
identity_indx = torch.eye(N, N, device=is_pos.device)
|
||||
remain_indx = torch.zeros(N, M - N, device=is_pos.device)
|
||||
identity_indx = torch.cat((identity_indx, remain_indx), dim=1)
|
||||
is_pos = is_pos - identity_indx
|
||||
else:
|
||||
is_pos = is_pos - torch.eye(N, N, device=is_pos.device)
|
||||
|
||||
is_neg = targets.view(N, 1).expand(N, M).ne(all_targets.view(M, 1).expand(M, N).t())
|
||||
|
||||
s_p = dist_mat * is_pos
|
||||
s_n = dist_mat * is_neg
|
||||
|
||||
alpha_p = torch.clamp_min(-s_p.detach() + 1 + self._m, min=0.)
|
||||
alpha_n = torch.clamp_min(s_n.detach() + self._m, min=0.)
|
||||
delta_p = 1 - self._m
|
||||
delta_n = self._m
|
||||
|
||||
logit_p = - self._s * alpha_p * (s_p - delta_p)
|
||||
logit_n = self._s * alpha_n * (s_n - delta_n)
|
||||
|
||||
loss = nn.functional.softplus(torch.logsumexp(logit_p, dim=1) + torch.logsumexp(logit_n, dim=1)).mean()
|
||||
|
||||
return loss * self._scale
|
||||
|
|
@ -58,5 +58,11 @@ class CrossEntropyLoss(object):
|
|||
targets *= smooth_param / (self._num_classes - 1)
|
||||
targets.scatter_(1, gt_classes.data.unsqueeze(1), (1 - smooth_param))
|
||||
|
||||
loss = (-targets * log_probs).mean(0).sum()
|
||||
loss = (-targets * log_probs).sum(dim=1)
|
||||
|
||||
with torch.no_grad():
|
||||
non_zero_cnt = max(loss.nonzero().size(0), 1)
|
||||
|
||||
loss = loss.sum() / non_zero_cnt
|
||||
|
||||
return loss * self._scale
|
||||
|
|
|
|||
|
|
@ -0,0 +1,241 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
# based on:
|
||||
# https://github.com/Andrew-Brown1/Smooth_AP/blob/master/src/Smooth_AP_loss.py
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from fastreid.utils import comm
|
||||
from fastreid.modeling.losses.utils import concat_all_gather
|
||||
|
||||
|
||||
def sigmoid(tensor, temp=1.0):
|
||||
""" temperature controlled sigmoid
|
||||
takes as input a torch tensor (tensor) and passes it through a sigmoid, controlled by temperature: temp
|
||||
"""
|
||||
exponent = -tensor / temp
|
||||
# clamp the input tensor for stability
|
||||
exponent = torch.clamp(exponent, min=-50, max=50)
|
||||
y = 1.0 / (1.0 + torch.exp(exponent))
|
||||
return y
|
||||
|
||||
|
||||
class SmoothAP(object):
|
||||
r"""PyTorch implementation of the Smooth-AP loss.
|
||||
implementation of the Smooth-AP loss. Takes as input the mini-batch of CNN-produced feature embeddings and returns
|
||||
the value of the Smooth-AP loss. The mini-batch must be formed of a defined number of classes. Each class must
|
||||
have the same number of instances represented in the mini-batch and must be ordered sequentially by class.
|
||||
e.g. the labels for a mini-batch with batch size 9, and 3 represented classes (A,B,C) must look like:
|
||||
labels = ( A, A, A, B, B, B, C, C, C)
|
||||
(the order of the classes however does not matter)
|
||||
For each instance in the mini-batch, the loss computes the Smooth-AP when it is used as the query and the rest of the
|
||||
mini-batch is used as the retrieval set. The positive set is formed of the other instances in the batch from the
|
||||
same class. The loss returns the average Smooth-AP across all instances in the mini-batch.
|
||||
Args:
|
||||
anneal : float
|
||||
the temperature of the sigmoid that is used to smooth the ranking function. A low value of the temperature
|
||||
results in a steep sigmoid, that tightly approximates the heaviside step function in the ranking function.
|
||||
batch_size : int
|
||||
the batch size being used during training.
|
||||
num_id : int
|
||||
the number of different classes that are represented in the batch.
|
||||
feat_dims : int
|
||||
the dimension of the input feature embeddings
|
||||
Shape:
|
||||
- Input (preds): (batch_size, feat_dims) (must be a cuda torch float tensor)
|
||||
- Output: scalar
|
||||
Examples::
|
||||
>>> loss = SmoothAP(0.01, 60, 6, 256)
|
||||
>>> input = torch.randn(60, 256, requires_grad=True).cuda()
|
||||
>>> output = loss(input)
|
||||
>>> output.backward()
|
||||
"""
|
||||
|
||||
def __init__(self, cfg):
|
||||
r"""
|
||||
Parameters
|
||||
----------
|
||||
cfg: (cfgNode)
|
||||
|
||||
anneal : float
|
||||
the temperature of the sigmoid that is used to smooth the ranking function
|
||||
batch_size : int
|
||||
the batch size being used
|
||||
num_id : int
|
||||
the number of different classes that are represented in the batch
|
||||
feat_dims : int
|
||||
the dimension of the input feature embeddings
|
||||
"""
|
||||
|
||||
self.anneal = 0.01
|
||||
self.num_id = cfg.SOLVER.IMS_PER_BATCH // cfg.DATALOADER.NUM_INSTANCE
|
||||
# self.num_id = 6
|
||||
|
||||
def __call__(self, embedding, targets):
|
||||
"""Forward pass for all input predictions: preds - (batch_size x feat_dims) """
|
||||
|
||||
# ------ differentiable ranking of all retrieval set ------
|
||||
embedding = F.normalize(embedding, dim=1)
|
||||
|
||||
feat_dim = embedding.size(1)
|
||||
|
||||
# For distributed training, gather all features from different process.
|
||||
if comm.get_world_size() > 1:
|
||||
all_embedding = concat_all_gather(embedding)
|
||||
all_targets = concat_all_gather(targets)
|
||||
else:
|
||||
all_embedding = embedding
|
||||
all_targets = targets
|
||||
|
||||
sim_dist = torch.matmul(embedding, all_embedding.t())
|
||||
N, M = sim_dist.size()
|
||||
|
||||
# Compute the mask which ignores the relevance score of the query to itself
|
||||
mask_indx = 1.0 - torch.eye(M, device=sim_dist.device)
|
||||
mask_indx = mask_indx.unsqueeze(dim=0).repeat(N, 1, 1) # (N, M, M)
|
||||
|
||||
# sim_dist -> N, 1, M -> N, M, N
|
||||
sim_dist_repeat = sim_dist.unsqueeze(dim=1).repeat(1, M, 1) # (N, M, M)
|
||||
# sim_dist_repeat_t = sim_dist.t().unsqueeze(dim=1).repeat(1, N, 1) # (N, N, M)
|
||||
|
||||
# Compute the difference matrix
|
||||
sim_diff = sim_dist_repeat - sim_dist_repeat.permute(0, 2, 1) # (N, M, M)
|
||||
|
||||
# Pass through the sigmoid
|
||||
sim_sg = sigmoid(sim_diff, temp=self.anneal) * mask_indx
|
||||
|
||||
# Compute all the rankings
|
||||
sim_all_rk = torch.sum(sim_sg, dim=-1) + 1 # (N, N)
|
||||
|
||||
pos_mask = targets.view(N, 1).expand(N, M).eq(all_targets.view(M, 1).expand(M, N).t()).float() # (N, M)
|
||||
|
||||
pos_mask_repeat = pos_mask.unsqueeze(1).repeat(1, M, 1) # (N, M, M)
|
||||
|
||||
# Compute positive rankings
|
||||
pos_sim_sg = sim_sg * pos_mask_repeat
|
||||
sim_pos_rk = torch.sum(pos_sim_sg, dim=-1) + 1 # (N, N)
|
||||
|
||||
# sum the values of the Smooth-AP for all instances in the mini-batch
|
||||
ap = 0
|
||||
group = N // self.num_id
|
||||
for ind in range(self.num_id):
|
||||
pos_divide = torch.sum(
|
||||
sim_pos_rk[(ind * group):((ind + 1) * group), (ind * group):((ind + 1) * group)] / (sim_all_rk[(ind * group):((ind + 1) * group), (ind * group):((ind + 1) * group)]))
|
||||
ap += pos_divide / torch.sum(pos_mask[ind*group]) / N
|
||||
return 1 - ap
|
||||
|
||||
|
||||
class SmoothAP_old(torch.nn.Module):
|
||||
"""PyTorch implementation of the Smooth-AP loss.
|
||||
implementation of the Smooth-AP loss. Takes as input the mini-batch of CNN-produced feature embeddings and returns
|
||||
the value of the Smooth-AP loss. The mini-batch must be formed of a defined number of classes. Each class must
|
||||
have the same number of instances represented in the mini-batch and must be ordered sequentially by class.
|
||||
e.g. the labels for a mini-batch with batch size 9, and 3 represented classes (A,B,C) must look like:
|
||||
labels = ( A, A, A, B, B, B, C, C, C)
|
||||
(the order of the classes however does not matter)
|
||||
For each instance in the mini-batch, the loss computes the Smooth-AP when it is used as the query and the rest of the
|
||||
mini-batch is used as the retrieval set. The positive set is formed of the other instances in the batch from the
|
||||
same class. The loss returns the average Smooth-AP across all instances in the mini-batch.
|
||||
Args:
|
||||
anneal : float
|
||||
the temperature of the sigmoid that is used to smooth the ranking function. A low value of the temperature
|
||||
results in a steep sigmoid, that tightly approximates the heaviside step function in the ranking function.
|
||||
batch_size : int
|
||||
the batch size being used during training.
|
||||
num_id : int
|
||||
the number of different classes that are represented in the batch.
|
||||
feat_dims : int
|
||||
the dimension of the input feature embeddings
|
||||
Shape:
|
||||
- Input (preds): (batch_size, feat_dims) (must be a cuda torch float tensor)
|
||||
- Output: scalar
|
||||
Examples::
|
||||
>>> loss = SmoothAP(0.01, 60, 6, 256)
|
||||
>>> input = torch.randn(60, 256, requires_grad=True).cuda()
|
||||
>>> output = loss(input)
|
||||
>>> output.backward()
|
||||
"""
|
||||
|
||||
def __init__(self, anneal, batch_size, num_id, feat_dims):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
anneal : float
|
||||
the temperature of the sigmoid that is used to smooth the ranking function
|
||||
batch_size : int
|
||||
the batch size being used
|
||||
num_id : int
|
||||
the number of different classes that are represented in the batch
|
||||
feat_dims : int
|
||||
the dimension of the input feature embeddings
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
assert(batch_size%num_id==0)
|
||||
|
||||
self.anneal = anneal
|
||||
self.batch_size = batch_size
|
||||
self.num_id = num_id
|
||||
self.feat_dims = feat_dims
|
||||
|
||||
def forward(self, preds):
|
||||
"""Forward pass for all input predictions: preds - (batch_size x feat_dims) """
|
||||
|
||||
preds = F.normalize(preds, dim=1)
|
||||
# ------ differentiable ranking of all retrieval set ------
|
||||
# compute the mask which ignores the relevance score of the query to itself
|
||||
mask = 1.0 - torch.eye(self.batch_size)
|
||||
mask = mask.unsqueeze(dim=0).repeat(self.batch_size, 1, 1)
|
||||
# compute the relevance scores via cosine similarity of the CNN-produced embedding vectors
|
||||
sim_all = torch.mm(preds, preds.t())
|
||||
sim_all_repeat = sim_all.unsqueeze(dim=1).repeat(1, self.batch_size, 1)
|
||||
# compute the difference matrix
|
||||
sim_diff = sim_all_repeat - sim_all_repeat.permute(0, 2, 1)
|
||||
# pass through the sigmoid
|
||||
sim_sg = sigmoid(sim_diff, temp=self.anneal) * mask
|
||||
# compute the rankings
|
||||
sim_all_rk = torch.sum(sim_sg, dim=-1) + 1
|
||||
|
||||
# ------ differentiable ranking of only positive set in retrieval set ------
|
||||
# compute the mask which only gives non-zero weights to the positive set
|
||||
xs = preds.view(self.num_id, int(self.batch_size / self.num_id), self.feat_dims)
|
||||
pos_mask = 1.0 - torch.eye(int(self.batch_size / self.num_id))
|
||||
pos_mask = pos_mask.unsqueeze(dim=0).unsqueeze(dim=0).repeat(self.num_id, int(self.batch_size / self.num_id), 1, 1)
|
||||
# compute the relevance scores
|
||||
sim_pos = torch.bmm(xs, xs.permute(0, 2, 1))
|
||||
sim_pos_repeat = sim_pos.unsqueeze(dim=2).repeat(1, 1, int(self.batch_size / self.num_id), 1)
|
||||
# compute the difference matrix
|
||||
sim_pos_diff = sim_pos_repeat - sim_pos_repeat.permute(0, 1, 3, 2)
|
||||
# pass through the sigmoid
|
||||
sim_pos_sg = sigmoid(sim_pos_diff, temp=self.anneal) * pos_mask
|
||||
# compute the rankings of the positive set
|
||||
sim_pos_rk = torch.sum(sim_pos_sg, dim=-1) + 1
|
||||
|
||||
# sum the values of the Smooth-AP for all instances in the mini-batch
|
||||
ap = torch.zeros(1)
|
||||
group = int(self.batch_size / self.num_id)
|
||||
for ind in range(self.num_id):
|
||||
pos_divide = torch.sum(sim_pos_rk[ind] / (sim_all_rk[(ind * group):((ind + 1) * group), (ind * group):((ind + 1) * group)]))
|
||||
ap = ap + ((pos_divide / group) / self.batch_size)
|
||||
|
||||
return 1-ap
|
||||
|
||||
if __name__ == '__main__':
|
||||
loss1 = SmoothAP(0.01)
|
||||
loss2 = SmoothAP_old(0.01, 60, 6, 256)
|
||||
|
||||
inputs = torch.randn(60, 256, requires_grad=True)
|
||||
targets = []
|
||||
for i in range(6):
|
||||
targets.extend([i]*10)
|
||||
targets = torch.LongTensor(targets)
|
||||
|
||||
output1 = loss1(inputs, targets)
|
||||
output2 = loss2(inputs)
|
||||
|
||||
print(torch.sum(output1 - output2))
|
||||
|
|
@ -8,51 +8,7 @@ import torch
|
|||
import torch.nn.functional as F
|
||||
|
||||
from fastreid.utils import comm
|
||||
|
||||
|
||||
# utils
|
||||
@torch.no_grad()
|
||||
def concat_all_gather(tensor):
|
||||
"""
|
||||
Performs all_gather operation on the provided tensors.
|
||||
*** Warning ***: torch.distributed.all_gather has no gradient.
|
||||
"""
|
||||
tensors_gather = [torch.ones_like(tensor)
|
||||
for _ in range(torch.distributed.get_world_size())]
|
||||
torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
|
||||
|
||||
output = torch.cat(tensors_gather, dim=0)
|
||||
return output
|
||||
|
||||
|
||||
def normalize(x, axis=-1):
|
||||
"""Normalizing to unit length along the specified dimension.
|
||||
Args:
|
||||
x: pytorch Variable
|
||||
Returns:
|
||||
x: pytorch Variable, same shape as input
|
||||
"""
|
||||
x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)
|
||||
return x
|
||||
|
||||
|
||||
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
|
||||
from .utils import concat_all_gather, euclidean_dist, normalize
|
||||
|
||||
|
||||
def softmax_weights(dist, mask):
|
||||
|
|
@ -174,42 +130,3 @@ class TripletLoss(object):
|
|||
if loss == float('Inf'): loss = F.margin_ranking_loss(dist_an, dist_ap, y, margin=0.3)
|
||||
|
||||
return loss * self._scale
|
||||
|
||||
|
||||
class CircleLoss(object):
|
||||
def __init__(self, cfg):
|
||||
self._scale = cfg.MODEL.LOSSES.CIRCLE.SCALE
|
||||
|
||||
self.m = cfg.MODEL.LOSSES.CIRCLE.MARGIN
|
||||
self.s = cfg.MODEL.LOSSES.CIRCLE.ALPHA
|
||||
|
||||
def __call__(self, embedding, targets):
|
||||
embedding = F.normalize(embedding, dim=1)
|
||||
|
||||
if comm.get_world_size() > 1:
|
||||
all_embedding = concat_all_gather(embedding)
|
||||
all_targets = concat_all_gather(targets)
|
||||
else:
|
||||
all_embedding = embedding
|
||||
all_targets = targets
|
||||
|
||||
dist_mat = torch.matmul(embedding, all_embedding.t())
|
||||
|
||||
N, M = dist_mat.size()
|
||||
is_pos = targets.view(N, 1).expand(N, M).eq(all_targets.view(M, 1).expand(M, N).t())
|
||||
is_neg = targets.view(N, 1).expand(N, M).ne(all_targets.view(M, 1).expand(M, N).t())
|
||||
|
||||
s_p = dist_mat[is_pos].contiguous().view(N, -1)
|
||||
s_n = dist_mat[is_neg].contiguous().view(N, -1)
|
||||
|
||||
alpha_p = F.relu(-s_p.detach() + 1 + self.m)
|
||||
alpha_n = F.relu(s_n.detach() + self.m)
|
||||
delta_p = 1 - self.m
|
||||
delta_n = self.m
|
||||
|
||||
logit_p = - self.s * alpha_p * (s_p - delta_p)
|
||||
logit_n = self.s * alpha_n * (s_n - delta_n)
|
||||
|
||||
loss = F.softplus(torch.logsumexp(logit_p, dim=1) + torch.logsumexp(logit_n, dim=1)).mean()
|
||||
|
||||
return loss * self._scale
|
||||
|
|
@ -0,0 +1,51 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def concat_all_gather(tensor):
|
||||
"""
|
||||
Performs all_gather operation on the provided tensors.
|
||||
*** Warning ***: torch.distributed.all_gather has no gradient.
|
||||
"""
|
||||
tensors_gather = [torch.ones_like(tensor)
|
||||
for _ in range(torch.distributed.get_world_size())]
|
||||
torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
|
||||
|
||||
output = torch.cat(tensors_gather, dim=0)
|
||||
return output
|
||||
|
||||
|
||||
def normalize(x, axis=-1):
|
||||
"""Normalizing to unit length along the specified dimension.
|
||||
Args:
|
||||
x: pytorch Variable
|
||||
Returns:
|
||||
x: pytorch Variable, same shape as input
|
||||
"""
|
||||
x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)
|
||||
return x
|
||||
|
||||
|
||||
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
|
||||
|
|
@ -28,7 +28,8 @@ class Baseline(nn.Module):
|
|||
|
||||
# head
|
||||
pool_type = cfg.MODEL.HEADS.POOL_LAYER
|
||||
if pool_type == 'avgpool': pool_layer = FastGlobalAvgPool2d()
|
||||
if pool_type == 'fastavgpool': pool_layer = FastGlobalAvgPool2d()
|
||||
elif pool_type == 'avgpool': pool_layer = nn.AdaptiveAvgPool2d(1)
|
||||
elif pool_type == 'maxpool': pool_layer = nn.AdaptiveMaxPool2d(1)
|
||||
elif pool_type == 'gempool': pool_layer = GeneralizedMeanPoolingP()
|
||||
elif pool_type == "avgmaxpool": pool_layer = AdaptiveAvgMaxPool2d()
|
||||
|
|
@ -66,8 +67,10 @@ class Baseline(nn.Module):
|
|||
"""
|
||||
Normalize and batch the input images.
|
||||
"""
|
||||
images = batched_inputs["images"].to(self.device)
|
||||
# images = batched_inputs
|
||||
if isinstance(batched_inputs, dict):
|
||||
images = batched_inputs["images"].to(self.device)
|
||||
elif isinstance(batched_inputs, torch.Tensor):
|
||||
images = batched_inputs.to(self.device)
|
||||
images.sub_(self.pixel_mean).div_(self.pixel_std)
|
||||
return images
|
||||
|
||||
|
|
@ -89,4 +92,7 @@ class Baseline(nn.Module):
|
|||
if "TripletLoss" in loss_names:
|
||||
loss_dict['loss_triplet'] = TripletLoss(self._cfg)(pred_features, gt_labels)
|
||||
|
||||
if "CircleLoss" in loss_names:
|
||||
loss_dict['loss_circle'] = CircleLoss(self._cfg)(pred_features, gt_labels)
|
||||
|
||||
return loss_dict
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
# based on:
|
||||
# https://github.com/pytorch/contrib/blob/master/torchcontrib/optim/swa.py
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
# based on
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import math
|
||||
|
|
@ -35,5 +35,3 @@ def weights_init_classifier(m):
|
|||
nn.init.normal_(m.weight, std=0.001)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0.0)
|
||||
elif classname.find("Arcface") != -1 or classname.find("Circle") != -1:
|
||||
nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5))
|
||||
|
|
|
|||
|
|
@ -3,11 +3,10 @@ MODEL:
|
|||
|
||||
BACKBONE:
|
||||
NAME: "build_resnet_backbone"
|
||||
DEPTH: 50
|
||||
DEPTH: "50x"
|
||||
NORM: "BN"
|
||||
LAST_STRIDE: 1
|
||||
WITH_IBN: True
|
||||
PRETRAIN_PATH: "/export/home/lxy/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar"
|
||||
|
||||
HEADS:
|
||||
NAME: "DSRHead"
|
||||
|
|
|
|||
|
|
@ -1,38 +1,29 @@
|
|||
# Deployment
|
||||
# Model Deployment
|
||||
|
||||
This directory contains:
|
||||
|
||||
1. A script that converts a fastreid model to Caffe format.
|
||||
1. The scripts that convert a fastreid model to Caffe/ONNX/TRT format.
|
||||
|
||||
2. An exmpale that loads a R50 baseline model in Caffe and run inference.
|
||||
2. The exmpales that load a R50 baseline model in Caffe/ONNX/TRT and run inference.
|
||||
|
||||
## Tutorial
|
||||
|
||||
### Caffe Convert
|
||||
|
||||
<details>
|
||||
<summary>step-to-step pipeline for caffe convert</summary>
|
||||
|
||||
This is a tiny example for converting fastreid-baseline in `meta_arch` to Caffe model, if you want to convert more complex architecture, you need to customize more things.
|
||||
|
||||
1. Change `preprocess_image` in `fastreid/modeling/meta_arch/baseline.py` as below
|
||||
|
||||
```python
|
||||
def preprocess_image(self, batched_inputs):
|
||||
"""
|
||||
Normalize and batch the input images.
|
||||
"""
|
||||
# images = [x["images"] for x in batched_inputs]
|
||||
# images = batched_inputs["images"]
|
||||
images = batched_inputs
|
||||
images.sub_(self.pixel_mean).div_(self.pixel_std)
|
||||
return images
|
||||
```
|
||||
|
||||
2. Run `caffe_export.py` to get the converted Caffe model,
|
||||
1. Run `caffe_export.py` to get the converted Caffe model,
|
||||
|
||||
```bash
|
||||
python caffe_export.py --config-file "/export/home/lxy/fast-reid/logs/market1501/bagtricks_R50/config.yaml" --name "baseline_R50" --output "logs/caffe_model" --opts MODEL.WEIGHTS "/export/home/lxy/fast-reid/logs/market1501/bagtricks_R50/model_final.pth"
|
||||
python caffe_export.py --config-file root-path/market1501/bagtricks_R50/config.yml --name "baseline_R50" --output outputs/caffe_model --opts MODEL.WEIGHTS root-path/logs/market1501/bagtricks_R50/model_final.pth
|
||||
```
|
||||
|
||||
then you can check the Caffe model and prototxt in `logs/caffe_model`.
|
||||
then you can check the Caffe model and prototxt in `outputs/caffe_model`.
|
||||
|
||||
3. Change `prototxt` following next three steps:
|
||||
2. Change `prototxt` following next three steps:
|
||||
|
||||
1) Edit `max_pooling` in `baseline_R50.prototxt` like this
|
||||
|
||||
|
|
@ -67,7 +58,7 @@ This is a tiny example for converting fastreid-baseline in `meta_arch` to Caffe
|
|||
}
|
||||
```
|
||||
|
||||
3) Change the last layer `top` name to `output`
|
||||
3) Change the last layer `top` name to `output`
|
||||
|
||||
```prototxt
|
||||
layer {
|
||||
|
|
@ -81,22 +72,89 @@ This is a tiny example for converting fastreid-baseline in `meta_arch` to Caffe
|
|||
}
|
||||
```
|
||||
|
||||
4. (optional) You can open [Netscope](https://ethereon.github.io/netscope/quickstart.html), then enter you network `prototxt` to visualize the network.
|
||||
3. (optional) You can open [Netscope](https://ethereon.github.io/netscope/quickstart.html), then enter you network `prototxt` to visualize the network.
|
||||
|
||||
5. Run `caffe_inference.py` to save Caffe model features with input images
|
||||
4. Run `caffe_inference.py` to save Caffe model features with input images
|
||||
|
||||
```bash
|
||||
python caffe_inference.py --model-def "logs/caffe_model/baseline_R50.prototxt" \
|
||||
--model-weights "logs/caffe_model/baseline_R50.caffemodel" \
|
||||
--input \
|
||||
'/export/home/DATA/Market-1501-v15.09.15/bounding_box_test/1182_c5s3_015240_04.jpg' \
|
||||
'/export/home/DATA/Market-1501-v15.09.15/bounding_box_test/1182_c6s3_038217_01.jpg' \
|
||||
'/export/home/DATA/Market-1501-v15.09.15/bounding_box_test/1183_c5s3_006943_05.jpg' \
|
||||
--output "caffe_R34_output"
|
||||
python caffe_inference.py --model-def outputs/caffe_model/baseline_R50.prototxt \
|
||||
--model-weights outputs/caffe_model/baseline_R50.caffemodel \
|
||||
--input test_data/*.jpg --output caffe_output
|
||||
```
|
||||
|
||||
6. Run `demo/demo.py` to get fastreid model features with the same input images, then compute the cosine similarity of difference model features to verify if you convert Caffe model successfully.
|
||||
6. Run `demo/demo.py` to get fastreid model features with the same input images, then verify that Caffe and PyTorch are computing the same value for the network.
|
||||
|
||||
```python
|
||||
np.testing.assert_allclose(torch_out, ort_out, rtol=1e-3, atol=1e-6)
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
### ONNX Convert
|
||||
|
||||
<details>
|
||||
<summary>step-to-step pipeline for onnx convert</summary>
|
||||
|
||||
This is a tiny example for converting fastreid-baseline in `meta_arch` to ONNX model. ONNX supports most operators in pytorch as far as I know and if some operators are not supported by ONNX, you need to customize these.
|
||||
|
||||
1. Run `onnx_export.py` to get the converted ONNX model,
|
||||
|
||||
```bash
|
||||
python onnx_export.py --config-file root-path/bagtricks_R50/config.yml --name "baseline_R50" --output outputs/onnx_model --opts MODEL.WEIGHTS root-path/logs/market1501/bagtricks_R50/model_final.pth
|
||||
```
|
||||
|
||||
then you can check the ONNX model in `outputs/onnx_model`.
|
||||
|
||||
2. (optional) You can use [Netron](https://github.com/lutzroeder/netron) to visualize the network.
|
||||
|
||||
3. Run `onnx_inference.py` to save ONNX model features with input images
|
||||
|
||||
```bash
|
||||
python onnx_inference.py --model-path outputs/onnx_model/baseline_R50.onnx \
|
||||
--input test_data/*.jpg --output onnx_output
|
||||
```
|
||||
|
||||
4. Run `demo/demo.py` to get fastreid model features with the same input images, then verify that ONNX Runtime and PyTorch are computing the same value for the network.
|
||||
|
||||
```python
|
||||
np.testing.assert_allclose(torch_out, ort_out, rtol=1e-3, atol=1e-6)
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
### TensorRT Convert
|
||||
|
||||
<details>
|
||||
<summary>step-to-step pipeline for trt convert</summary>
|
||||
|
||||
This is a tiny example for converting fastreid-baseline in `meta_arch` to TRT model. We use [tiny-tensorrt](https://github.com/zerollzeng/tiny-tensorrt) which is a simple and easy-to-use nvidia TensorRt warpper, to get the model converted to tensorRT.
|
||||
|
||||
First you need to convert the pytorch model to ONNX format following [ONNX Convert](https://github.com/JDAI-CV/fast-reid#fastreid), and you need to remember your `output` name. Then you can convert ONNX model to TensorRT following instructions below.
|
||||
|
||||
1. Run command line below to get the converted TRT model from ONNX model,
|
||||
|
||||
```bash
|
||||
|
||||
python trt_export.py --name "baseline_R50" --output outputs/trt_model --onnx-model outputs/onnx_model/baseline.onnx --heighi 256 --width 128
|
||||
```
|
||||
|
||||
then you can check the TRT model in `outputs/trt_model`.
|
||||
|
||||
2. Run `trt_inference.py` to save TRT model features with input images
|
||||
|
||||
```bash
|
||||
python onnx_inference.py --model-path outputs/trt_model/baseline.engine \
|
||||
--input test_data/*.jpg --output trt_output --output-name trt_model_outputname
|
||||
```
|
||||
|
||||
3. Run `demo/demo.py` to get fastreid model features with the same input images, then verify that TensorRT and PyTorch are computing the same value for the network.
|
||||
|
||||
```python
|
||||
np.testing.assert_allclose(torch_out, ort_out, rtol=1e-3, atol=1e-6)
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## Acknowledgements
|
||||
|
||||
Thank to [CPFLAME](https://github.com/CPFLAME), [gcong18](https://github.com/gcong18), [YuxiangJohn](https://github.com/YuxiangJohn) and [wiggin66](https://github.com/wiggin66) at JDAI Model Acceleration Group for help in PyTorch to Caffe model converting.
|
||||
Thank to [CPFLAME](https://github.com/CPFLAME), [gcong18](https://github.com/gcong18), [YuxiangJohn](https://github.com/YuxiangJohn) and [wiggin66](https://github.com/wiggin66) at JDAI Model Acceleration Group for help in PyTorch model converting.
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
|
@ -15,6 +15,9 @@ from fastreid.config import get_cfg
|
|||
from fastreid.modeling.meta_arch import build_model
|
||||
from fastreid.utils.file_io import PathManager
|
||||
from fastreid.utils.checkpoint import Checkpointer
|
||||
from fastreid.utils.logger import setup_logger
|
||||
|
||||
logger = setup_logger(name='caffe_export')
|
||||
|
||||
|
||||
def setup_cfg(args):
|
||||
|
|
@ -64,10 +67,12 @@ if __name__ == '__main__':
|
|||
model = build_model(cfg)
|
||||
Checkpointer(model).load(cfg.MODEL.WEIGHTS)
|
||||
model.eval()
|
||||
print(model)
|
||||
logger.info(model)
|
||||
|
||||
inputs = torch.randn(1, 3, cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1]).cuda()
|
||||
inputs = torch.randn(1, 3, cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1])
|
||||
PathManager.mkdirs(args.output)
|
||||
pytorch_to_caffe.trans_net(model, inputs, args.name)
|
||||
pytorch_to_caffe.save_prototxt(f"{args.output}/{args.name}.prototxt")
|
||||
pytorch_to_caffe.save_caffemodel(f"{args.output}/{args.name}.caffemodel")
|
||||
|
||||
logger.info(f"Export caffe model in {args.output} sucessfully!")
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: liaoxingyu5@jd.com
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import caffe
|
||||
|
|
@ -43,7 +43,7 @@ def get_parser():
|
|||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=384,
|
||||
default=256,
|
||||
help="height of image"
|
||||
)
|
||||
parser.add_argument(
|
||||
|
|
|
|||
|
|
@ -1,48 +0,0 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import torch
|
||||
sys.path.append('../..')
|
||||
from fastreid.config import get_cfg
|
||||
from fastreid.engine import default_argument_parser, default_setup
|
||||
from fastreid.modeling.meta_arch import build_model
|
||||
from fastreid.export.tensorflow_export import export_tf_reid_model
|
||||
from fastreid.export.tf_modeling import TfMetaArch
|
||||
|
||||
|
||||
def setup(args):
|
||||
"""
|
||||
Create configs and perform basic setups.
|
||||
"""
|
||||
cfg = get_cfg()
|
||||
# cfg.merge_from_file(args.config_file)
|
||||
cfg.merge_from_list(args.opts)
|
||||
cfg.freeze()
|
||||
default_setup(cfg, args)
|
||||
return cfg
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = default_argument_parser().parse_args()
|
||||
print("Command Line Args:", args)
|
||||
cfg = setup(args)
|
||||
cfg.defrost()
|
||||
cfg.MODEL.BACKBONE.NAME = "build_resnet_backbone"
|
||||
cfg.MODEL.BACKBONE.DEPTH = 50
|
||||
cfg.MODEL.BACKBONE.LAST_STRIDE = 1
|
||||
# If use IBN block in backbone
|
||||
cfg.MODEL.BACKBONE.WITH_IBN = False
|
||||
cfg.MODEL.BACKBONE.PRETRAIN = False
|
||||
|
||||
from torchvision.models import resnet50
|
||||
# model = TfMetaArch(cfg)
|
||||
model = resnet50(pretrained=False)
|
||||
# model.load_params_wo_fc(torch.load('logs/bjstation/res50_baseline_v0.4/ckpts/model_epoch80.pth'))
|
||||
model.eval()
|
||||
dummy_inputs = torch.randn(1, 3, 256, 128)
|
||||
export_tf_reid_model(model, dummy_inputs, 'reid_tf.pb')
|
||||
|
|
@ -0,0 +1,146 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import io
|
||||
import sys
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
from onnxsim import simplify
|
||||
from torch.onnx import OperatorExportTypes
|
||||
|
||||
sys.path.append('../../')
|
||||
|
||||
from fastreid.config import get_cfg
|
||||
from fastreid.modeling.meta_arch import build_model
|
||||
from fastreid.utils.file_io import PathManager
|
||||
from fastreid.utils.checkpoint import Checkpointer
|
||||
from fastreid.utils.logger import setup_logger
|
||||
|
||||
logger = setup_logger(name='onnx_export')
|
||||
|
||||
|
||||
def setup_cfg(args):
|
||||
cfg = get_cfg()
|
||||
cfg.merge_from_file(args.config_file)
|
||||
cfg.merge_from_list(args.opts)
|
||||
cfg.freeze()
|
||||
return cfg
|
||||
|
||||
|
||||
def get_parser():
|
||||
parser = argparse.ArgumentParser(description="Convert Pytorch to ONNX model")
|
||||
|
||||
parser.add_argument(
|
||||
"--config-file",
|
||||
metavar="FILE",
|
||||
help="path to config file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--name",
|
||||
default="baseline",
|
||||
help="name for converted model"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default='onnx_model',
|
||||
help='path to save converted onnx model'
|
||||
)
|
||||
parser.add_argument(
|
||||
"--opts",
|
||||
help="Modify config options using the command-line 'KEY VALUE' pairs",
|
||||
default=[],
|
||||
nargs=argparse.REMAINDER,
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def remove_initializer_from_input(model):
|
||||
if model.ir_version < 4:
|
||||
print(
|
||||
'Model with ir_version below 4 requires to include initilizer in graph input'
|
||||
)
|
||||
return
|
||||
|
||||
inputs = model.graph.input
|
||||
name_to_input = {}
|
||||
for input in inputs:
|
||||
name_to_input[input.name] = input
|
||||
|
||||
for initializer in model.graph.initializer:
|
||||
if initializer.name in name_to_input:
|
||||
inputs.remove(name_to_input[initializer.name])
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def export_onnx_model(model, inputs):
|
||||
"""
|
||||
Trace and export a model to onnx format.
|
||||
Args:
|
||||
model (nn.Module):
|
||||
inputs (torch.Tensor): the model will be called by `model(*inputs)`
|
||||
Returns:
|
||||
an onnx model
|
||||
"""
|
||||
assert isinstance(model, torch.nn.Module)
|
||||
|
||||
# make sure all modules are in eval mode, onnx may change the training state
|
||||
# of the module if the states are not consistent
|
||||
def _check_eval(module):
|
||||
assert not module.training
|
||||
|
||||
model.apply(_check_eval)
|
||||
|
||||
# Export the model to ONNX
|
||||
with torch.no_grad():
|
||||
with io.BytesIO() as f:
|
||||
torch.onnx.export(
|
||||
model,
|
||||
inputs,
|
||||
f,
|
||||
operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK,
|
||||
# verbose=True, # NOTE: uncomment this for debugging
|
||||
# export_params=True,
|
||||
)
|
||||
onnx_model = onnx.load_from_string(f.getvalue())
|
||||
|
||||
# Apply ONNX's Optimization
|
||||
all_passes = onnx.optimizer.get_available_passes()
|
||||
passes = ["extract_constant_to_initializer", "eliminate_unused_initializer", "fuse_bn_into_conv"]
|
||||
assert all(p in all_passes for p in passes)
|
||||
onnx_model = onnx.optimizer.optimize(onnx_model, passes)
|
||||
return onnx_model
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = get_parser().parse_args()
|
||||
cfg = setup_cfg(args)
|
||||
|
||||
cfg.defrost()
|
||||
cfg.MODEL.BACKBONE.PRETRAIN = False
|
||||
if cfg.MODEL.HEADS.POOL_LAYER == 'fastavgpool':
|
||||
cfg.MODEL.HEADS.POOL_LAYER = 'avgpool'
|
||||
model = build_model(cfg)
|
||||
Checkpointer(model).load(cfg.MODEL.WEIGHTS)
|
||||
model.eval()
|
||||
logger.info(model)
|
||||
|
||||
inputs = torch.randn(1, 3, cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1])
|
||||
onnx_model = export_onnx_model(model, inputs)
|
||||
|
||||
model_simp, check = simplify(onnx_model)
|
||||
|
||||
model_simp = remove_initializer_from_input(model_simp)
|
||||
|
||||
assert check, "Simplified ONNX model could not be validated"
|
||||
|
||||
PathManager.mkdirs(args.output)
|
||||
|
||||
onnx.save_model(model_simp, f"{args.output}/{args.name}.onnx")
|
||||
|
||||
logger.info(f"Export onnx model in {args.output} successfully!")
|
||||
|
|
@ -0,0 +1,85 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnxruntime
|
||||
import tqdm
|
||||
|
||||
|
||||
def get_parser():
|
||||
parser = argparse.ArgumentParser(description="onnx model inference")
|
||||
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
default="onnx_model/baseline.onnx",
|
||||
help="onnx model path"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input",
|
||||
nargs="+",
|
||||
help="A list of space separated input images; "
|
||||
"or a single glob pattern such as 'directory/*.jpg'",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default='onnx_output',
|
||||
help='path to save converted caffe model'
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=256,
|
||||
help="height of image"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=128,
|
||||
help="width of image"
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def preprocess(image_path, image_height, image_width):
|
||||
original_image = cv2.imread(image_path)
|
||||
# the model expects RGB inputs
|
||||
original_image = original_image[:, :, ::-1]
|
||||
|
||||
# Apply pre-processing to image.
|
||||
img = cv2.resize(original_image, (image_width, image_height), interpolation=cv2.INTER_CUBIC)
|
||||
img = img.astype("float32").transpose(2, 0, 1)[np.newaxis] # (1, 3, h, w)
|
||||
return img
|
||||
|
||||
|
||||
def normalize(nparray, order=2, axis=-1):
|
||||
"""Normalize a N-D numpy array along the specified axis."""
|
||||
norm = np.linalg.norm(nparray, ord=order, axis=axis, keepdims=True)
|
||||
return nparray / (norm + np.finfo(np.float32).eps)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = get_parser().parse_args()
|
||||
|
||||
ort_sess = onnxruntime.InferenceSession(args.model_path)
|
||||
|
||||
input_name = ort_sess.get_inputs()[0].name
|
||||
|
||||
if not os.path.exists(args.output): os.makedirs(args.output)
|
||||
|
||||
if args.input:
|
||||
if os.path.isdir(args.input[0]):
|
||||
args.input = glob.glob(os.path.expanduser(args.input[0]))
|
||||
assert args.input, "The input path(s) was not found"
|
||||
for path in tqdm.tqdm(args.input):
|
||||
image = preprocess(path, args.height, args.width)
|
||||
feat = ort_sess.run(None, {input_name: image})[0]
|
||||
feat = normalize(feat, axis=1)
|
||||
np.save(os.path.join(args.output, path.replace('.jpg', '.npy').split('/')[-1]), feat)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
|
||||
python caffe_inference.py --model-def "logs/caffe_R34/baseline_R34.prototxt" \
|
||||
--model-weights "logs/caffe_R34/baseline_R34.caffemodel" \
|
||||
--height 256 --width 128 \
|
||||
--input \
|
||||
'/export/home/DATA/Market-1501-v15.09.15/bounding_box_test/1182_c5s3_015240_04.jpg' \
|
||||
'/export/home/DATA/Market-1501-v15.09.15/bounding_box_test/1182_c6s3_038217_01.jpg' \
|
||||
'/export/home/DATA/Market-1501-v15.09.15/bounding_box_test/1183_c5s3_006943_05.jpg' \
|
||||
'/export/home/DATA/DukeMTMC-reID/bounding_box_train/0728_c4_f0161265.jpg' \
|
||||
--output "caffe_R34_output"
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 2.2 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.1 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.1 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.8 KiB |
|
|
@ -0,0 +1,82 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import numpy as np
|
||||
import sys
|
||||
|
||||
sys.path.append('../../')
|
||||
sys.path.append("/export/home/lxy/runtimelib-tensorrt-tiny/build")
|
||||
|
||||
import pytrt
|
||||
from fastreid.utils.logger import setup_logger
|
||||
from fastreid.utils.file_io import PathManager
|
||||
|
||||
|
||||
logger = setup_logger(name='trt_export')
|
||||
|
||||
|
||||
def get_parser():
|
||||
parser = argparse.ArgumentParser(description="Convert ONNX to TRT model")
|
||||
|
||||
parser.add_argument(
|
||||
"--name",
|
||||
default="baseline",
|
||||
help="name for converted model"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default='outputs/trt_model',
|
||||
help='path to save converted trt model'
|
||||
)
|
||||
parser.add_argument(
|
||||
"--onnx-model",
|
||||
default='outputs/onnx_model/baseline.onnx',
|
||||
help='path to onnx model'
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=256,
|
||||
help="height of image"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=128,
|
||||
help="width of image"
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def export_trt_model(onnxModel, engineFile, input_numpy_array):
|
||||
r"""
|
||||
Export a model to trt format.
|
||||
"""
|
||||
|
||||
trt = pytrt.Trt()
|
||||
|
||||
customOutput = []
|
||||
maxBatchSize = 1
|
||||
calibratorData = []
|
||||
mode = 2
|
||||
trt.CreateEngine(onnxModel, engineFile, customOutput, maxBatchSize, mode, calibratorData)
|
||||
trt.DoInference(input_numpy_array) # slightly different from c++
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = get_parser().parse_args()
|
||||
|
||||
inputs = np.zeros(shape=(32, args.height, args.width, 3))
|
||||
onnxModel = args.onnx_model
|
||||
engineFile = os.path.join(args.output, args.name+'.engine')
|
||||
|
||||
PathManager.mkdirs(args.output)
|
||||
export_trt_model(onnxModel, engineFile, inputs)
|
||||
|
||||
logger.info(f"Export trt model in {args.output} successfully!")
|
||||
|
|
@ -0,0 +1,99 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: xingyu liao
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import sys
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
# import tqdm
|
||||
|
||||
sys.path.append("/export/home/lxy/runtimelib-tensorrt-tiny/build")
|
||||
|
||||
import pytrt
|
||||
|
||||
|
||||
def get_parser():
|
||||
parser = argparse.ArgumentParser(description="trt model inference")
|
||||
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
default="outputs/trt_model/baseline.engine",
|
||||
help="trt model path"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input",
|
||||
nargs="+",
|
||||
help="A list of space separated input images; "
|
||||
"or a single glob pattern such as 'directory/*.jpg'",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default="trt_output",
|
||||
help="path to save trt model inference results"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-name",
|
||||
help="tensorRT model output name"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=256,
|
||||
help="height of image"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=128,
|
||||
help="width of image"
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def preprocess(image_path, image_height, image_width):
|
||||
original_image = cv2.imread(image_path)
|
||||
# the model expects RGB inputs
|
||||
original_image = original_image[:, :, ::-1]
|
||||
|
||||
# Apply pre-processing to image.
|
||||
img = cv2.resize(original_image, (image_width, image_height), interpolation=cv2.INTER_CUBIC)
|
||||
img = img.astype("float32").transpose(2, 0, 1)[np.newaxis] # (1, 3, h, w)
|
||||
return img
|
||||
|
||||
|
||||
def normalize(nparray, order=2, axis=-1):
|
||||
"""Normalize a N-D numpy array along the specified axis."""
|
||||
norm = np.linalg.norm(nparray, ord=order, axis=axis, keepdims=True)
|
||||
return nparray / (norm + np.finfo(np.float32).eps)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = get_parser().parse_args()
|
||||
|
||||
trt = pytrt.Trt()
|
||||
|
||||
onnxModel = ""
|
||||
engineFile = args.model_path
|
||||
customOutput = []
|
||||
maxBatchSize = 1
|
||||
calibratorData = []
|
||||
mode = 2
|
||||
trt.CreateEngine(onnxModel, engineFile, customOutput, maxBatchSize, mode, calibratorData)
|
||||
|
||||
if not os.path.exists(args.output): os.makedirs(args.output)
|
||||
|
||||
if args.input:
|
||||
if os.path.isdir(args.input[0]):
|
||||
args.input = glob.glob(os.path.expanduser(args.input[0]))
|
||||
assert args.input, "The input path(s) was not found"
|
||||
for path in args.input:
|
||||
input_numpy_array = preprocess(path, args.height, args.width)
|
||||
trt.DoInference(input_numpy_array)
|
||||
feat = trt.GetOutput(args.output_name)
|
||||
feat = normalize(feat, axis=1)
|
||||
np.save(os.path.join(args.output, path.replace('.jpg', '.npy').split('/')[-1]), feat)
|
||||
Loading…
Reference in New Issue