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

update code

pull/43/head
liaoxingyu 2019-08-20 09:36:47 +08:00
parent fd751ef197
commit aa4c8eee2c
25 changed files with 563 additions and 971 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
.gitignore vendored
View File

@ -2,6 +2,6 @@
__pycache__
.DS_Store
.vscode
csrc/eval_cylib/build/
csrc/eval_cylib/*.so
logs/
.ipynb_checkpoints
.ipynb_checkpoints

18
README.md
View File

@ -7,6 +7,7 @@ We support
- [x] end-to-end training and evaluation
- [ ] multi-GPU distributed training
- [ ] fast training speed with fp16
- [x] fast evaluation with cython
- [ ] support both image and video reid
- [x] multi-dataset training
- [x] cross-dataset evaluation
@ -15,7 +16,7 @@ We support
- [ ] high efficient backbone
- [ ] advanced training techniques
- [ ] various loss functions
- [ ] visualization tools
- [ ] tensorboard visualization
## Get Started
The designed architecture follows this guide [PyTorch-Project-Template](https://github.com/L1aoXingyu/PyTorch-Project-Template), you can check each folder's purpose by yourself.
@ -42,12 +43,15 @@ The designed architecture follows this guide [PyTorch-Project-Template](https://
bounding_box_test/
bounding_box_train/
```
5. Prepare pretrained model if you don't have
```python
from torchvision import models
models.resnet50(pretrained=True)
5. Prepare pretrained model.
If you use origin ResNet, you do not need to do anything. But if you want to use ResNet_ibn, you need to download pretrain model in [here](https://drive.google.com/open?id=1thS2B8UOSBi_cJX6zRy6YYRwz_nVFI_S). And then you can put it in `~/.cache/torch/checkpoints` or anywhere you like.
Then you should set this pretrain model path in `configs/softmax_triplet.yml`.
6. compile with cython to accelerate evalution
```bash
cd csrc/eval_cylib; make
```
Then it will automatically download model in `~/.cache/torch/checkpoints/`, you should set this path in `config/defaults.py` for all training or set in every single training config file in `configs/`.
## Train
Most of the configuration files that we provide, you can run this command for training market1501
@ -71,4 +75,4 @@ python3 tools/test.py --config_file='configs/softmax.yml' TEST.WEIGHT '/save/tra
| cfg | market1501 | dukemtmc |
| --- | -- | -- |
| softmax_triplet, size=(256, 128), batch_size=64(16 id x 4 imgs) | 93.9 (85.9) | training |
| softmax+triplet, size=(256, 128), batch_size=64(16 id x 4 imgs) | 93.9 (85.9) | 86.5 (75.9) |

15
config/defaults.py
View File

@ -16,11 +16,24 @@ from yacs.config import CfgNode as CN
_C = CN()
# -----------------------------------------------------------------------------
# MODEL
# -----------------------------------------------------------------------------
_C.MODEL = CN()
_C.MODEL.DEVICE = "cuda"
# Model backbone
_C.MODEL.BACKBONE = 'resnet50'
# Last stride for backbone
_C.MODEL.LAST_STRIDE = 1
# If use IBN block
_C.MODEL.IBN = False
# If use imagenet pretrain model
_C.MODEL.PRETRAIN = True
# Pretrain model path
_C.MODEL.PRETRAIN_PATH = ''
# Checkpoint for continuing training
_C.MODEL.CHECKPOINT = ''
#
# -----------------------------------------------------------------------------
# INPUT
# -----------------------------------------------------------------------------
@ -69,7 +82,7 @@ _C.DATALOADER.NUM_INSTANCE = 16
# Solver
# ---------------------------------------------------------------------------- #
_C.SOLVER = CN()
_C.SOLVER.OPTIMIZER_NAME = "Adam"
_C.SOLVER.OPT = "adam"
_C.SOLVER.MAX_EPOCHS = 50

3
configs/softmax_triplet.yml
View File

@ -1,6 +1,5 @@
MODEL:
BACKBONE: "resnet50"
PRETRAIN_PATH: '/export/home/lxy/.cache/torch/checkpoints/resnet50-19c8e357.pth'
INPUT:
@ -18,7 +17,7 @@ DATALOADER:
NUM_INSTANCE: 4
SOLVER:
OPTIMIZER_NAME: 'Adam'
OPT: 'adam'
MAX_EPOCHS: 150
BASE_LR: 0.00035
WEIGHT_DECAY: 0.0005

17
data/build.py
View File

@ -36,23 +36,24 @@ def get_data_bunch(cfg):
market_query_path = 'datasets/Market-1501-v15.09.15/query'
marker_gallery_path = 'datasets/Market-1501-v15.09.15/bounding_box_test'
duke_query_path = 'datasets/DukeMTMC-reID/query'
duke_gallery_path = 'datasets/DukeMTMC-reID/bounding_box_test'
train_img_names = list()
for d in cfg.DATASETS.NAMES:
if d == 'market1501':
train_img_names.extend(_process_dir(market_train_path))
elif d == 'duke':
train_img_names.extend(_process_dir(duke_train_path))
elif d == 'cuhk03':
train_img_names.extend(CUHK03().train)
else:
raise NameError(f'{d} is not available')
if d == 'market1501': train_img_names.extend(_process_dir(market_train_path))
elif d == 'duke': train_img_names.extend(_process_dir(duke_train_path))
elif d == 'cuhk03': train_img_names.extend(CUHK03().train)
else: raise NameError(f'{d} is not available')
train_names = [i[0] for i in train_img_names]
if cfg.DATASETS.TEST_NAMES == "market1501":
query_names = _process_dir(market_query_path)
gallery_names = _process_dir(marker_gallery_path)
elif cfg.DATASETS.TEST_NAMES == 'duke':
query_names = _process_dir(duke_query_path)
gallery_names = _process_dir(duke_gallery_path)
else:
print(f"not support {cfg.DATASETS.TEST_NAMES} test set")

79
engine/callbacks.py 100644
View File

@ -0,0 +1,79 @@
# encoding: utf-8
"""
@author: liaoxingyu
@contact: sherlockliao01@gmail.com
"""
from fastai.vision import *
import logging
from data.datasets.eval_reid import evaluate
__all__ = ['TrackValue', 'LRScheduler', 'TestModel']
@dataclass
class TrackValue(Callback):
logger: logging.Logger
total_iter: int
def on_epoch_end(self, epoch, smooth_loss, **kwargs):
self.logger.info(f'Epoch {epoch}[Iter {self.total_iter}], loss: {smooth_loss.item():.4f}')
class LRScheduler(LearnerCallback):
def __init__(self, learn, lr_sched):
super().__init__(learn)
self.lr_sched = lr_sched
def on_train_begin(self, **kwargs:Any):
self.opt = self.learn.opt
def on_epoch_begin(self, **kwargs:Any):
self.opt.lr = self.lr_sched.step()
class TestModel(LearnerCallback):
def __init__(self, learn: Learner, test_labels: Iterator, eval_period: int, num_query: int, logger: logging.Logger, norm=True):
super().__init__(learn)
self._test_dl = learn.data.test_dl
self._eval_period = eval_period
self._norm = norm
self._logger = logger
self._num_query = num_query
pids = []
camids = []
for i in test_labels:
pids.append(i[0])
camids.append(i[1])
self.q_pids = np.asarray(pids[:num_query])
self.q_camids = np.asarray(camids[:num_query])
self.g_pids = np.asarray(pids[num_query:])
self.g_camids = np.asarray(camids[num_query:])
def on_epoch_end(self, epoch, **kwargs: Any):
# test model performance
if (epoch + 1) % self._eval_period == 0:
self._logger.info('Testing ...')
feats, pids, camids = [], [], []
self.learn.model.eval()
with torch.no_grad():
for imgs, _ in self._test_dl:
feat = self.learn.model(imgs)
feats.append(feat)
feats = torch.cat(feats, dim=0)
if self._norm:
feats = F.normalize(feats, p=2, dim=1)
# query
qf = feats[:self._num_query]
# gallery
gf = feats[self._num_query:]
m, n = qf.shape[0], gf.shape[0]
distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
distmat.addmm_(1, -2, qf, gf.t())
distmat = to_np(distmat)
cmc, mAP = evaluate(distmat, self.q_pids, self.g_pids, self.q_camids, self.g_camids)
self._logger.info(f"Test Results - Epoch: {epoch+1}")
self._logger.info(f"mAP: {mAP:.1%}")
for r in [1, 5, 10]:
self._logger.info(f"CMC curve, Rank-{r:<3}:{cmc[r-1]:.1%}")
self.learn.save("model_{}".format(epoch))

10
engine/inference.py
View File

@ -9,7 +9,6 @@ import torch
import numpy as np
import torch.nn.functional as F
from data.datasets.eval_reid import evaluate
from data.datasets.eval_threshold import eval_roc
from fastai.torch_core import to_np
@ -32,6 +31,7 @@ def inference(
g_pids = np.asarray(pids[num_query:])
q_camids = np.asarray(camids[:num_query])
g_camids = np.asarray(camids[num_query:])
feats = []
model.eval()
for imgs, _ in data_bunch.test_dl:
@ -60,11 +60,3 @@ def inference(
logger.info("mAP: {:.1%}".format(mAP))
for r in [1, 5, 10]:
logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(r, cmc[r - 1]))
# Compute ROC and AUC
logger.info("Compute ROC Curve...")
fpr, tpr, fps, tps, p, n, thresholds = eval_roc(distmat, q_pids, g_pids, q_camids, g_camids, 0.1, 0.5)
logger.info("positive samples: {}, negative samples: {}".format(p, n))
for i, thresh in enumerate(thresholds):
logger.info("threshold: {:.2f}, FP: {:.0f}({:.3f}), TP: {:.0f}({:.3f})".
format(thresh, fps[i], fpr[i], tps[i], tpr[i]))

16
engine/interpreter.py
View File

@ -7,20 +7,20 @@
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from fastai.basic_data import *
from fastai.layers import *
from fastai.vision import *
class ReidInterpretation():
"Interpretation methods for reid models."
def __init__(self, learn, test_labels, num_q):
self.test_labels,self.num_q = test_labels,num_q
self.learn,self.test_labels,self.num_q = learn,test_labels,num_q
self.test_dl = learn.data.test_dl
self.model = learn.model
self.get_distmat()
def get_distmat(self):
self.model.eval()
feats = []
pids = []
camids = []
for p,c in self.test_labels:
@ -31,11 +31,7 @@ class ReidInterpretation():
self.q_camids = np.asarray(camids[:self.num_q])
self.g_camids = np.asarray(camids[self.num_q:])
for imgs, _ in self.test_dl:
with torch.no_grad():
feat = self.model(imgs)
feats.append(feat)
feats = torch.cat(feats, dim=0)
feats, _ = self.learn.get_preds(DatasetType.Test, activ=Lambda(lambda x:x))
feats = F.normalize(feats)
qf = feats[:self.num_q]
gf = feats[self.num_q:]

85
engine/trainer.py
View File

@ -8,79 +8,8 @@ import logging
import os
import matplotlib.pyplot as plt
import torch.nn.functional as F
from data.datasets.eval_reid import evaluate
from fastai.vision import *
@dataclass
class TrackValue(Callback):
logger: logging.Logger
total_iter: int
def on_epoch_end(self, epoch, smooth_loss, **kwargs):
self.logger.info(f'Epoch {epoch}[Iter {self.total_iter}], loss: {smooth_loss.item():.4f}')
@dataclass
class LRScheduler(Callback):
learn: Learner
lr_sched: Scheduler
def on_train_begin(self, **kwargs:Any):
self.opt = self.learn.opt
def on_epoch_begin(self, **kwargs:Any):
self.opt.lr = self.lr_sched.step()
class TestModel(LearnerCallback):
def __init__(self, learn: Learner, test_labels: Iterator, eval_period: int, num_query: int, logger: logging.Logger, norm=True):
super().__init__(learn)
self._test_dl = learn.data.test_dl
self._eval_period = eval_period
self._norm = norm
self._logger = logger
self._num_query = num_query
pids = []
camids = []
for i in test_labels:
pids.append(i[0])
camids.append(i[1])
self.q_pids = np.asarray(pids[:num_query])
self.q_camids = np.asarray(camids[:num_query])
self.g_pids = np.asarray(pids[num_query:])
self.g_camids = np.asarray(camids[num_query:])
def on_epoch_end(self, epoch, **kwargs: Any):
# test model performance
if (epoch + 1) % self._eval_period == 0:
self._logger.info('Testing ...')
feats, pids, camids = [], [], []
self.learn.model.eval()
with torch.no_grad():
for imgs, _ in self._test_dl:
feat = self.learn.model(imgs)
feats.append(feat)
feats = torch.cat(feats, dim=0)
if self._norm:
feats = F.normalize(feats, p=2, dim=1)
# query
qf = feats[:self._num_query]
# gallery
gf = feats[self._num_query:]
m, n = qf.shape[0], gf.shape[0]
distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
distmat.addmm_(1, -2, qf, gf.t())
distmat = to_np(distmat)
cmc, mAP = evaluate(distmat, self.q_pids, self.g_pids, self.q_camids, self.g_camids)
self._logger.info(f"Test Results - Epoch: {epoch+1}")
self._logger.info(f"mAP: {mAP:.1%}")
for r in [1, 5, 10]:
self._logger.info(f"CMC curve, Rank-{r:<3}:{cmc[r-1]:.1%}")
self.learn.save("model_{}".format(epoch))
from .callbacks import *
def do_train(
@ -116,4 +45,16 @@ def do_train(
callback_fns=cb_fns,
callbacks=[TrackValue(logger, total_iter)])
# continue training
if cfg.MODEL.CHECKPOINT is not '':
state = torch.load(cfg.MODEL.CHECKPOINT)
if set(state.keys()) == {'model', 'opt'}:
model_state = state['model']
learn.model.load_state_dict(model_state)
learn.create_opt(0, 0)
learn.opt.load_state_dict(state['opt'])
else:
learn.model.load_state_dict(state['model'])
logger.info(f'continue training from checkpoint {cfg.MODEL.CHECKPOINT}')
learn.fit(epochs, lr=cfg.SOLVER.BASE_LR, wd=cfg.SOLVER.WEIGHT_DECAY)

14
layers/triplet_loss.py
View File

@ -5,6 +5,7 @@
"""
import torch
from torch import nn
import torch.nn.functional as F
def normalize(x, axis=-1):
@ -62,12 +63,20 @@ def hard_example_mining(dist_mat, labels, return_inds=False):
# `dist_ap` means distance(anchor, positive)
# both `dist_ap` and `relative_p_inds` with shape [N, 1]
# pos_dist = dist_mat[is_pos].contiguous().view(N, -1)
# ap_weight = F.softmax(pos_dist, dim=1)
# dist_ap = torch.sum(ap_weight * pos_dist, dim=1)
dist_ap, relative_p_inds = torch.max(
dist_mat[is_pos].contiguous().view(N, -1), 1, keepdim=True)
# `dist_an` means distance(anchor, negative)
# both `dist_an` and `relative_n_inds` with shape [N, 1]
dist_an, relative_n_inds = torch.min(
dist_mat[is_neg].contiguous().view(N, -1), 1, keepdim=True)
# neg_dist = dist_mat[is_neg].contiguous().view(N, -1)
# an_weight = F.softmax(-neg_dist, dim=1)
# dist_an = torch.sum(an_weight * neg_dist, dim=1)
# shape [N]
dist_ap = dist_ap.squeeze(1)
dist_an = dist_an.squeeze(1)
@ -90,19 +99,20 @@ def hard_example_mining(dist_mat, labels, return_inds=False):
return dist_ap, dist_an
class TripletLoss(object):
class TripletLoss(nn.Module):
"""Modified from Tong Xiao's open-reid (https://github.com/Cysu/open-reid).
Related Triplet Loss theory can be found in paper 'In Defense of the Triplet
Loss for Person Re-Identification'."""
def __init__(self, margin=None):
super().__init__()
self.margin = margin
if margin is not None:
self.ranking_loss = nn.MarginRankingLoss(margin=margin)
else:
self.ranking_loss = nn.SoftMarginLoss()
def __call__(self, global_feat, labels, normalize_feature=False):
def forward(self, global_feat, labels, normalize_feature=False):
if normalize_feature:
global_feat = normalize(global_feat, axis=-1)
dist_mat = euclidean_dist(global_feat, global_feat)

3
modeling/__init__.py
View File

@ -8,5 +8,6 @@ from .baseline import Baseline
def build_model(cfg, num_classes):
model = Baseline(cfg.MODEL.BACKBONE, num_classes, cfg.MODEL.LAST_STRIDE, cfg.MODEL.PRETRAIN_PATH)
model = Baseline(cfg.MODEL.BACKBONE, num_classes, cfg.MODEL.LAST_STRIDE,
cfg.MODEL.IBN, cfg.MODEL.PRETRAIN, cfg.MODEL.PRETRAIN_PATH)
return model

98
modeling/backbones/resnet.py
View File

@ -8,17 +8,56 @@ import math
import torch
from torch import nn
from torch.utils import model_zoo
model_urls = {
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
}
model_layers = {
'resnet50': [3, 4, 6, 3],
'resnet101': [3, 4, 23, 3]
}
__all__ = ['ResNet']
__all__ = ['resnet50']
class IBN(nn.Module):
def __init__(self, planes):
super(IBN, self).__init__()
half1 = int(planes/2)
self.half = half1
half2 = planes - half1
self.IN = nn.InstanceNorm2d(half1, affine=True)
self.BN = nn.BatchNorm2d(half2)
def forward(self, x):
split = torch.split(x, self.half, 1)
out1 = self.IN(split[0].contiguous())
# out2 = self.BN(torch.cat(split[1:], dim=1).contiguous())
out2 = self.BN(split[1].contiguous())
out = torch.cat((out1, out2), 1)
return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
if ibn:
self.bn1 = IBN(planes)
else:
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
@ -52,21 +91,22 @@ class Bottleneck(nn.Module):
class ResNet(nn.Module):
def __init__(self, last_stride=2, block=Bottleneck, layers=[3, 4, 6, 3]):
self.inplanes = 64
super(ResNet, self).__init__()
def __init__(self, last_stride, ibn, block, layers):
scale = 64
self.inplanes = scale
super().__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer1 = self._make_layer(block, scale, layers[0], ibn=ibn)
self.layer2 = self._make_layer(block, scale*2, layers[1], stride=2, ibn=ibn)
self.layer3 = self._make_layer(block, scale*4, layers[2], stride=2, ibn=ibn)
self.layer4 = self._make_layer(
block, 512, layers[3], stride=last_stride)
block, scale*8, layers[3], stride=last_stride)
def _make_layer(self, block, planes, blocks, stride=1):
def _make_layer(self, block, planes, blocks, stride=1, ibn=False):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
@ -76,10 +116,12 @@ class ResNet(nn.Module):
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
if planes == 512:
ibn = False
layers.append(block(self.inplanes, planes, ibn, stride, downsample))
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes))
layers.append(block(self.inplanes, planes, ibn))
return nn.Sequential(*layers)
@ -96,12 +138,22 @@ class ResNet(nn.Module):
return x
def load_param(self, model_path):
param_dict = torch.load(model_path)
for i in param_dict:
if 'fc' in i:
continue
self.state_dict()[i].copy_(param_dict[i])
def load_pretrain(self, model_path=''):
if model_path == '':
state_dict = model_zoo.load_url(model_urls[self._model_name])
state_dict.pop('fc.weight')
state_dict.pop('fc.bias')
else:
state_dict = torch.load(model_path)['state_dict']
state_dict.pop('module.fc.weight')
state_dict.pop('module.fc.bias')
new_state_dict = {}
for k in state_dict:
new_k = '.'.join(k.split('.')[1:]) # remove module in name
if self.state_dict()[new_k].shape == state_dict[k].shape:
new_state_dict[new_k] = state_dict[k]
state_dict = new_state_dict
self.load_state_dict(state_dict, strict=False)
def random_init(self):
for m in self.modules():
@ -112,7 +164,7 @@ class ResNet(nn.Module):
m.weight.data.fill_(1)
m.bias.data.zero_()
def resnet50(last_stride, **kwargs):
model = ResNet(last_stride, block=Bottleneck, layers=[3,4,6,3])
return model
@classmethod
def from_name(cls, model_name, last_stride, ibn):
cls._model_name = model_name
return ResNet(last_stride, ibn=ibn, block=Bottleneck, layers=model_layers[model_name])

27
modeling/backbones/resnet_ibn_a.py
View File

@ -1,18 +1,10 @@
import math
import torch
import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo
from torch.utils import model_zoo
__all__ = ['ResNet_IBN', 'resnet50_ibn_a', 'resnet101_ibn_a',
'resnet152_ibn_a']
model_urls = {
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}
__all__ = ['ResNet_IBN', 'resnet50_ibn_a']
class IBN(nn.Module):
@ -144,6 +136,15 @@ class ResNet_IBN(nn.Module):
j = '.'.join(i.split('.')[1:]) # remove 'module' in state_dict
if self.state_dict()[j].shape == param_dict[i].shape:
self.state_dict()[j].copy_(param_dict[i])
def load_pretrain(self):
state_dict = model_zoo.load_url(model_urls[self._model_name])
@classmethod
def from_name(cls, model_name, last_stride):
cls._model_name = model_name
return ResNet_IBN(last_stride, Bottleneck_IBN, [3, 4, 6, 3])
def resnet50_ibn_a(last_stride, **kwargs):
@ -172,4 +173,4 @@ def resnet152_ibn_a(last_stride, pretrained=False, **kwargs):
model = ResNet_IBN(last_stride, Bottleneck_IBN, [3, 8, 36, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
return model
return model

24
modeling/baseline.py
View File

@ -35,19 +35,12 @@ def weights_init_classifier(m):
class Baseline(nn.Module):
in_planes = 2048
def __init__(self, backbone, num_classes, last_stride, model_path=None):
super(Baseline, self).__init__()
if backbone == 'resnet50':
self.base = resnet50(last_stride)
elif backbone == 'resnet50_ibn':
self.base = resnet50_ibn_a(last_stride)
else:
print(f'not support {backbone} backbone')
def __init__(self, backbone, num_classes, last_stride, ibn, pretrain=True, model_path=None):
super().__init__()
try: self.base = ResNet.from_name(backbone, last_stride, ibn)
except: print(f'not support {backbone} backbone')
try:
self.base.load_param(model_path)
except:
print("not load imagenet pretrained model!")
if pretrain: self.base.load_pretrain(model_path)
self.gap = nn.AdaptiveAvgPool2d(1)
self.num_classes = num_classes
@ -70,7 +63,6 @@ class Baseline(nn.Module):
return feat
def load_params_wo_fc(self, state_dict):
for i in state_dict:
if 'classifier' in i:
continue
self.state_dict()[i].copy_(state_dict[i])
state_dict.pop('classifier.weight')
res = self.load_state_dict(state_dict, strict=False)
assert str(res.missing_keys) == str(['classifier.weight',]), 'issue loading pretrained weights'

8
scripts/test_model.sh
View File

@ -2,8 +2,8 @@ gpu=0
CUDA_VISIBLE_DEVICES=$gpu python tools/test.py -cfg='configs/softmax_triplet.yml' \
MODEL.BACKBONE 'resnet50' \
INPUT.SIZE_TRAIN '(256, 128)' \
DATASETS.NAMES '("market1501","duke","beijing")' \
DATASETS.TEST_NAMES 'market1501' \
MODEL.IBN 'True' \
MODEL.PRETRAIN 'False' \
DATASETS.TEST_NAMES 'duke' \
OUTPUT_DIR 'logs/test' \
TEST.WEIGHT 'logs/market/bs64_light/models/model_149.pth'
TEST.WEIGHT 'logs/2019.8.16/market/resnet50_ibn_1_1/models/model_149.pth'

10
scripts/train_duke.sh 100644
View File

@ -0,0 +1,10 @@
gpu=0
CUDA_VISIBLE_DEVICES=$gpu python tools/train.py -cfg='configs/softmax_triplet.yml' \
DATASETS.NAMES '("duke",)' \
DATASETS.TEST_NAMES 'duke' \
MODEL.BACKBONE 'resnet50' \
MODEL.IBN 'False' \
INPUT.DO_LIGHTING 'False' \
SOLVER.OPT 'adam' \
OUTPUT_DIR 'logs/2019.8.19/duke/resnet'

7
scripts/train_market.sh
View File

@ -4,9 +4,12 @@ CUDA_VISIBLE_DEVICES=$gpu python tools/train.py -cfg='configs/softmax_triplet.ym
DATASETS.NAMES '("market1501",)' \
DATASETS.TEST_NAMES 'market1501' \
MODEL.BACKBONE 'resnet50' \
MODEL.PRETRAIN_PATH '/home/user01/.cache/torch/checkpoints/resnet50-19c8e357.pth' \
MODEL.IBN 'False' \
INPUT.DO_LIGHTING 'False' \
OUTPUT_DIR 'logs/2019.8.14/market/baseline'
SOLVER.OPT 'radam' \
OUTPUT_DIR 'logs/2019.8.17/market/resnet_radam_nowarmup'
# MODEL.PRETRAIN_PATH '/home/user01/.cache/torch/checkpoints/resnet50_ibn_a.pth.tar' \
# CUDA_VISIBLE_DEVICES=$gpu python tools/train.py -cfg='configs/softmax_triplet.yml' \
# DATASETS.NAMES '("market1501",)' \

8
solver/__init__.py 100644
View File

@ -0,0 +1,8 @@
# encoding: utf-8
"""
@author: liaoxingyu
@contact: sherlockliao01@gmail.com
"""
from .radam import *

210
solver/radam.py 100644
View File

@ -0,0 +1,210 @@
import math
import torch
from torch.optim.optimizer import Optimizer, required
__all__ = ['RAdam', 'PlainRAdam', 'AdamW']
class RAdam(Optimizer):
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
self.buffer = [[None, None, None] for ind in range(10)]
super(RAdam, self).__init__(params, defaults)
def __setstate__(self, state):
super(RAdam, self).__setstate__(state)
def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
grad = p.grad.data.float()
if grad.is_sparse:
raise RuntimeError('RAdam does not support sparse gradients')
p_data_fp32 = p.data.float()
state = self.state[p]
if len(state) == 0:
state['step'] = 0
state['exp_avg'] = torch.zeros_like(p_data_fp32)
state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)
else:
state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)
exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
beta1, beta2 = group['betas']
exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
exp_avg.mul_(beta1).add_(1 - beta1, grad)
state['step'] += 1
buffered = self.buffer[int(state['step'] % 10)]
if state['step'] == buffered[0]:
N_sma, step_size = buffered[1], buffered[2]
else:
buffered[0] = state['step']
beta2_t = beta2 ** state['step']
N_sma_max = 2 / (1 - beta2) - 1
N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
buffered[1] = N_sma
# more conservative since it's an approximated value
if N_sma >= 5:
step_size = group['lr'] * math.sqrt((1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / (N_sma_max - 2)) / (1 - beta1 ** state['step'])
else:
step_size = group['lr'] / (1 - beta1 ** state['step'])
buffered[2] = step_size
if group['weight_decay'] != 0:
p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)
# more conservative since it's an approximated value
if N_sma >= 5:
denom = exp_avg_sq.sqrt().add_(group['eps'])
p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
else:
p_data_fp32.add_(-step_size, exp_avg)
p.data.copy_(p_data_fp32)
return loss
class PlainRAdam(Optimizer):
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
super(RAdam, self).__init__(params, defaults)
def __setstate__(self, state):
super(RAdam, self).__setstate__(state)
def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
grad = p.grad.data.float()
if grad.is_sparse:
raise RuntimeError('RAdam does not support sparse gradients')
p_data_fp32 = p.data.float()
state = self.state[p]
if len(state) == 0:
state['step'] = 0
state['exp_avg'] = torch.zeros_like(p_data_fp32)
state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)
else:
state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)
exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
beta1, beta2 = group['betas']
exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
exp_avg.mul_(beta1).add_(1 - beta1, grad)
state['step'] += 1
beta2_t = beta2 ** state['step']
N_sma_max = 2 / (1 - beta2) - 1
N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
if group['weight_decay'] != 0:
p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)
# more conservative since it's an approximated value
if N_sma >= 5:
step_size = group['lr'] * math.sqrt((1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / (N_sma_max - 2)) / (1 - beta1 ** state['step'])
denom = exp_avg_sq.sqrt().add_(group['eps'])
p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
else:
step_size = group['lr'] / (1 - beta1 ** state['step'])
p_data_fp32.add_(-step_size, exp_avg)
p.data.copy_(p_data_fp32)
return loss
class AdamW(Optimizer):
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, warmup = 0):
defaults = dict(lr=lr, betas=betas, eps=eps,
weight_decay=weight_decay, amsgrad=amsgrad, use_variance=True, warmup = warmup)
super(AdamW, self).__init__(params, defaults)
def __setstate__(self, state):
super(AdamW, self).__setstate__(state)
def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
grad = p.grad.data.float()
if grad.is_sparse:
raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
p_data_fp32 = p.data.float()
state = self.state[p]
if len(state) == 0:
state['step'] = 0
state['exp_avg'] = torch.zeros_like(p_data_fp32)
state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)
else:
state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)
exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
beta1, beta2 = group['betas']
state['step'] += 1
exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
exp_avg.mul_(beta1).add_(1 - beta1, grad)
denom = exp_avg_sq.sqrt().add_(group['eps'])
bias_correction1 = 1 - beta1 ** state['step']
bias_correction2 = 1 - beta2 ** state['step']
if group['warmup'] > state['step']:
scheduled_lr = 1e-8 + state['step'] * group['lr'] / group['warmup']
else:
scheduled_lr = group['lr']
step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1
if group['weight_decay'] != 0:
p_data_fp32.add_(-group['weight_decay'] * scheduled_lr, p_data_fp32)
p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
p.data.copy_(p_data_fp32)
return loss

22
tests/dataset_test.py
View File

@ -1,22 +0,0 @@
# encoding: utf-8
"""
@author: liaoxingyu
@contact: sherlockliao01@gmail.com
"""
import sys
from fastai.vision import *
sys.path.append('.')
from data import get_data_bunch
from config import cfg
if __name__ == '__main__':
# cfg.INPUT.SIZE_TRAIN = (384, 128)
data, label, num_q = get_data_bunch(cfg)
# def get_ex(): return open_image('datasets/beijingStation/query/000245_c10s2_1561732033722.000000.jpg')
# im = get_ex()
print(data.train_ds[0])
print(data.test_ds[0])
from ipdb import set_trace; set_trace()
# im.apply_tfms(crop_pad(size=(300, 300)))

23
tests/interp_test.py 100644
View File

@ -0,0 +1,23 @@
import torch
from fastai.vision import *
from fastai.basic_data import *
from fastai.layers import *
import sys
sys.path.append('.')
from engine.interpreter import ReidInterpretation
from data import get_data_bunch
from modeling import build_model
from config import cfg
cfg.DATASETS.NAMES = ('market1501',)
cfg.DATASETS.TEST_NAMES = 'market1501'
cfg.MODEL.BACKBONE = 'resnet50'
data_bunch, test_labels, num_query = get_data_bunch(cfg)
model = build_model(cfg, 10)
model.load_params_wo_fc(torch.load('logs/2019.8.14/market/baseline/models/model_149.pth')['model'])
learn = Learner(data_bunch, model)
feats, _ = learn.get_preds(DatasetType.Test, activ=Lambda(lambda x: x))

25
tests/model_test.py 100644
View File

@ -0,0 +1,25 @@
import sys
import unittest
import torch
from torch import nn
import sys
sys.path.append('.')
from modeling.backbones import *
from config import cfg
class MyTestCase(unittest.TestCase):
def test_model(self):
net1 = ResNet.from_name('resnet50', 1, True)
for i in net1.named_parameters():
print(i[0])
net2 = resnet50_ibn_a(1)
# print('*'*10)
# for i in net2.named_parameters():
# print(i[0])
if __name__ == '__main__':
unittest.main()

10
tools/train.py
View File

@ -17,10 +17,10 @@ from engine.trainer import do_train
from fastai.vision import *
from layers import make_loss
from modeling import build_model
from solver import *
from utils.logger import setup_logger
def train(cfg):
# prepare dataset
data_bunch, test_labels, num_query = get_data_bunch(cfg)
@ -28,9 +28,11 @@ def train(cfg):
# prepare model
model = build_model(cfg, data_bunch.c)
opt_fns = partial(getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME))
if cfg.SOLVER.OPT == 'adam': opt_fns = partial(torch.optim.Adam)
elif cfg.SOLVER.OPT == 'sgd': opt_fns = partial(torch.optim.SGD, momentum=0.9)
else: raise NameError(f'optimizer {cfg.SOLVER.OPT} not support')
def warmup_multistep(start: float, end: float, pct: float):
def lr_multistep(start: float, end: float, pct: float):
warmup_factor = 1
gamma = cfg.SOLVER.GAMMA
milestones = [1.0 * s / cfg.SOLVER.MAX_EPOCHS for s in cfg.SOLVER.STEPS]
@ -40,7 +42,7 @@ def train(cfg):
warmup_factor = cfg.SOLVER.WARMUP_FACTOR * (1 - alpha) + alpha
return start * warmup_factor * gamma ** bisect_right(milestones, pct)
lr_sched = Scheduler(cfg.SOLVER.BASE_LR, cfg.SOLVER.MAX_EPOCHS, warmup_multistep)
lr_sched = Scheduler(cfg.SOLVER.BASE_LR, cfg.SOLVER.MAX_EPOCHS, lr_multistep)
loss_func = make_loss(cfg)

8
utils/logger.py
View File

@ -10,14 +10,6 @@ import sys
import logging
def mkdir_if_missing(dir_path):
try:
os.makedirs(dir_path)
except OSError as e:
if e.errno != errno.EEXIST:
raise
def setup_logger(name, save_dir, distributed_rank):
logger = logging.getLogger(name)
logger.setLevel(logging.DEBUG)

790
vis_data.ipynb
View File

File diff suppressed because one or more lines are too long