def make_test_loader_M_C_D(root, name):
query_dataset = Market1501(root + '/query',
transform=test_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
query_flip_dataset = Market1501(root + '/query',
transform=test_flip_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
query_loader = DataLoader(query_dataset,
batch_size=batch_test,
num_workers=num_workers)
query_flip_loader = DataLoader(query_flip_dataset,
batch_size=batch_test,
num_workers=num_workers)
test_dataset = Market1501(root + '/bounding_box_test',
transform=test_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
test_flip_dataset = Market1501(root + '/bounding_box_test',
transform=test_flip_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
test_loader = DataLoader(test_dataset,
batch_size=batch_test,
num_workers=num_workers)
test_flip_loader = DataLoader(test_flip_dataset,
batch_size=batch_test,
num_workers=num_workers)
return query_loader, query_flip_loader, test_loader, test_flip_loader
def read_and_prepare_data():
dataset = Market1501(FLAGS.inputs_path, num_validation_y=0.1, seed=1234)
X_train, Y_train, _ = dataset.read_train()
X_train = np.array(X_train)
Y_train = np.array(Y_train)
X_val, Y_val, _ = dataset.read_validation()
X_val = np.array(X_val)
Y_val = np.array(Y_val)
return X_train, Y_train, X_val, Y_val
def __init__(self, split, dataloader):
print("[*] Loading Market1501")
market = Market1501('gt_bbox', **dataloader)
print("[*] Loading CUHK03")
cuhk = CUHK03('labeled', **dataloader)
print("[*] Loading MOT")
mot = MOTreIDWrapper(split, dataloader)
self.dataset = ConcatDataset([market, cuhk, mot])
def train():
## data
logger.info('creating dataloader')
dataset = Market1501('./dataset/Market-1501-v15.09.15/bounding_box_train',
is_train=True)
num_classes = dataset.get_num_classes()
sampler = BalancedSampler(dataset, 16, 4)
dl = DataLoader(dataset, batch_sampler=sampler, num_workers=8)
## network and loss
logger.info('setup model and loss')
sphereloss = SphereLoss(1024, num_classes)
sphereloss.cuda()
net = Network_D()
net = nn.DataParallel(net)
net.train()
net.cuda()
## optimizer
logger.info('creating optimizer')
params = list(net.parameters())
params += list(sphereloss.parameters())
optim = torch.optim.Adam(params, lr=1e-3)
## training
logger.info('start training')
t_start = time.time()
loss_it = []
for ep in range(150):
optim, lrs = lr_scheduler(ep, optim)
for it, (imgs, lbs, ids) in enumerate(dl):
imgs = imgs.cuda()
lbs = lbs.cuda()
embs = net(imgs)
loss = sphereloss(embs, lbs)
optim.zero_grad()
loss.backward()
optim.step()
loss_it.append(loss.detach().cpu().numpy())
if it % 10 == 0 and it != 0:
t_end = time.time()
t_interval = t_end - t_start
log_loss = sum(loss_it) / len(loss_it)
msg = 'epoch: {}, iter: {}, loss: {:4f}, lr: {}, time: {:4f}'.format(
ep, it, log_loss, lrs, t_interval)
logger.info(msg)
loss_it = []
t_start = t_end
## save model
torch.save(net.module.state_dict(), './res/model_final.pkl')
logger.info('\nTraining done, model saved to {}\n\n'.format(
'./res/model_final.pkl'))
def setUp(self):
self.batch_id = 4
self.batch_image = 16
self.data_source = Market1501(root + '/bounding_box_train',
transform=ToTensor())
self.sampler = RandomIdSampler(self.data_source,
batch_image=self.batch_image)
self.data_loader = DataLoader(self.data_source,
sampler=self.sampler,
batch_size=self.batch_id *
self.batch_image)
def test_test(self):
test = Market1501(root + '/bounding_box_test')
self.assertEqual(15913, len(test.imgs))
self.assertEqual(15913, len(test.ids))
self.assertEqual(15913, len(test.cameras))
self.assertEqual(751, len(test._id2label))
self.assertTrue(test.imgs[0].endswith('0000_c1s1_000151_01.jpg'))
self.assertEqual(0, test.ids[0])
self.assertEqual(1, test.cameras[0])
self.assertEqual(0, test._id2label[0])
img, target = test[0]
self.assertEqual(0, target)
self.assertTrue(test.imgs[2798].endswith('0001_c1s1_001051_03.jpg'))
self.assertEqual(1, test.ids[2798])
self.assertEqual(1, test.cameras[2798])
self.assertEqual(1, test._id2label[1])
img, target = test[2798]
self.assertEqual(1, target)
def test_query(self):
query = Market1501(root + '/query')
self.assertEqual(3368, len(query.imgs))
self.assertEqual(3368, len(query.ids))
self.assertEqual(3368, len(query.cameras))
self.assertEqual(750, len(query._id2label))
self.assertTrue(query.imgs[0].endswith('0001_c1s1_001051_00.jpg'))
self.assertEqual(1, query.ids[0])
self.assertEqual(1, query.cameras[0])
self.assertEqual(0, query._id2label[1])
_, target = query[0]
self.assertEqual(0, target)
self.assertTrue(query.imgs[6].endswith('0003_c1s6_015971_00.jpg'))
self.assertEqual(3, query.ids[6])
self.assertEqual(1, query.cameras[6])
self.assertEqual(1, query._id2label[3])
_, target = query[6]
self.assertEqual(1, target)
def test_train(self):
train = Market1501(root + '/bounding_box_train')
self.assertEqual(12936, len(train.imgs))
self.assertEqual(12936, len(train.ids))
self.assertEqual(12936, len(train.cameras))
self.assertEqual(751, len(train._id2label))
self.assertTrue(train.imgs[0].endswith('0002_c1s1_000451_03.jpg'))
self.assertEqual(2, train.ids[0])
self.assertEqual(1, train.cameras[0])
self.assertEqual(0, train._id2label[2])
_, target = train[0]
self.assertEqual(0, target)
self.assertTrue(train.imgs[46].endswith('0007_c1s6_028546_01.jpg'))
self.assertEqual(7, train.ids[46])
self.assertEqual(1, train.cameras[46])
self.assertEqual(1, train._id2label[7])
_, target = train[46]
self.assertEqual(1, target)
def test_camera(self):
self.assertEqual(1, Market1501.camera('-1_c1s1_000401_03.jpg'))
def test_id(self):
self.assertEqual(-1, Market1501.id('-1_c1s1_000401_03.jpg'))
from market1501 import Market1501
from transform import Train_Transform, Val_Transform
from torch.utils.data import DataLoader
from model import Train_Model
from tripletloss import TripletLoss
from cross_entropy_smooth import CrossEntropySmooth
from optimizer import Make_Optimizer, Warmup
import os
import time
if torch.cuda.is_available():
print("Support GPU!")
device = torch.device('cuda')
else:
device = torch.device('cpu')
market = Market1501(root='./')
train_transform = Train_Transform(True)
val_transform = Val_Transform()
train_sampler = RandomIdentitySampler(market.train, 16, 4)
train_dataset = ImageDataset(dataset=market.train, transform=train_transform)
val_dataset = ImageDataset(dataset=market.test + market.query,
transform=val_transform)
train_dataloader = DataLoader(train_dataset, 64, False, train_sampler)
val_dataloader = DataLoader(val_dataset, 128, False)
Model = Train_Model().to(device)
IDloss = CrossEntropySmooth(market.num_train_id)
optimizer = Make_Optimizer(Model, 3.5e-5)
tripletloss = TripletLoss(0.3)
warmup = Warmup(optimizer)
EPOCH = 120
for epoch in range(EPOCH):
def run():
batch_id = 16
batch_image = 4
batch_test = 32
train_transform = transforms.Compose([
transforms.Resize((384, 128)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
train_dataset = Market1501(args.root + '/bounding_box_train', transform=train_transform)
train_loader = dataloader.DataLoader(train_dataset,
sampler=RandomIdSampler(train_dataset, batch_image=batch_image),
batch_size=batch_id * batch_image,
num_workers=args.workers)
test_transform = transforms.Compose([
transforms.Resize((384, 128)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
test_flip_transform = transforms.Compose([
transforms.Resize((384, 128)),
functional.hflip,
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
query_dataset = Market1501(args.root + '/query', transform=test_transform)
query_flip_dataset = Market1501(args.root + '/query', transform=test_flip_transform)
query_loader = dataloader.DataLoader(query_dataset, batch_size=batch_test, num_workers=args.workers)
query_flip_loader = dataloader.DataLoader(query_flip_dataset, batch_size=batch_test, num_workers=args.workers)
test_dataset = Market1501(args.root + '/bounding_box_test', transform=test_transform)
test_flip_dataset = Market1501(args.root + '/bounding_box_test', transform=test_flip_transform)
test_loader = dataloader.DataLoader(test_dataset, batch_size=batch_test, num_workers=args.workers)
test_flip_loader = dataloader.DataLoader(test_flip_dataset, batch_size=batch_test, num_workers=args.workers)
mgn = MGN(num_classes=len(train_dataset.unique_ids)).to(DEVICE)
cross_entropy_loss = nn.CrossEntropyLoss()
triplet_semihard_loss = TripletSemihardLoss(margin=1.2)
optimizer = optim.SGD(mgn.parameters(), lr=1e-2, momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [40, 60], gamma=0.1)
epochs = 80
for epoch in range(epochs):
mgn.train()
scheduler.step()
running_loss = 0.0
for i, data in enumerate(train_loader):
inputs, labels = data
inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
optimizer.zero_grad()
outputs = mgn(inputs)
losses = [triplet_semihard_loss(output, labels) for output in outputs[1]] + \
[cross_entropy_loss(output, labels) for output in outputs[2]]
loss = sum(losses) / len(losses)
loss.backward()
optimizer.step()
running_loss += loss.item()
print('%d/%d - %d/%d - loss: %f' % (epoch + 1, epochs, i, len(train_loader), loss.item()))
print('epoch: %d/%d - loss: %f' % (epoch + 1, epochs, running_loss / len(train_loader)))
if epoch % 10 == 9:
mgn.eval()
query = np.concatenate([mgn(inputs.to(DEVICE))[0].detach().cpu().numpy()
for inputs, _ in query_loader])
query_flip = np.concatenate([mgn(inputs.to(DEVICE))[0].detach().cpu().numpy()
for inputs, _ in query_flip_loader])
test = np.concatenate([mgn(inputs.to(DEVICE))[0].detach().cpu().numpy()
for inputs, _ in test_loader])
test_flip = np.concatenate([mgn(inputs.to(DEVICE))[0].detach().cpu().numpy()
for inputs, _ in test_flip_loader])
dist = cdist((query + query_flip) / 2., (test + test_flip) / 2.)
# dist = cdist(normalize(query + query_flip), normalize(test + test_flip))
r = cmc(dist, query_dataset.ids, test_dataset.ids, query_dataset.cameras, test_dataset.cameras,
separate_camera_set=False,
single_gallery_shot=False,
first_match_break=True)
m_ap = mean_ap(dist, query_dataset.ids, test_dataset.ids, query_dataset.cameras, test_dataset.cameras)
print('epoch[%d]: mAP=%f, r@1=%f, r@3=%f, r@5=%f, r@10=%f' % (epoch + 1, m_ap, r[0], r[2], r[4], r[9]))
def test_camera(self):
self.assertEqual(1, Market1501.camera('-1_c1s1_000401_03.jpg'))
def test_id(self):
self.assertEqual(-1, Market1501.id('-1_c1s1_000401_03.jpg'))
def train():
## data
P, K = 16, 8
batchsize = P * K
logger.info('creating dataloader')
dataset = Market1501('./dataset/Market-1501-v15.09.15/bounding_box_train',
is_train=True)
num_classes = dataset.get_num_classes()
sampler = BalancedSampler(dataset, P, K)
dl = DataLoader(dataset, batch_sampler=sampler, num_workers=8)
## network and loss
logger.info('setup model and loss')
bottleneck_loss = BottleneckLoss(2048, num_classes)
bottleneck_loss.cuda()
bottleneck_loss.train()
net = Embeddor()
net.cuda()
net.train()
# net = nn.DataParallel(net)
## optimizer
logger.info('creating optimizer')
lr = 0.1
momentum = 0.9
params = list(net.parameters())
params += list(bottleneck_loss.parameters())
optim = torch.optim.SGD(params, lr=lr, momentum=momentum)
## training
logger.info('start training')
n_epochs = 140
t_start = time.time()
loss_it = []
for ep in range(n_epochs):
optim, lrs = lr_scheduler(ep, optim)
for it, (imgs, lbs, _) in enumerate(dl):
imgs = imgs.cuda()
lbs = lbs.cuda()
optim.zero_grad()
embs_org, embs_sft = net(imgs)
loss_org = bottleneck_loss(embs_org, lbs)
loss_sft = bottleneck_loss(embs_sft, lbs)
loss = loss_org + loss_sft
loss.backward()
optim.step()
loss = loss.cpu().item()
loss_it.append(loss)
# print logging message
t_end = time.time()
t_interval = t_end - t_start
log_loss = sum(loss_it) / len(loss_it)
msg = ', '.join(
['epoch: {}', 'loss: {:.4f}', 'lr: {}',
'time: {:.4f}']).format(ep, log_loss, lrs, t_interval)
logger.info(msg)
loss_it = []
t_start = t_end
## save model
if hasattr(net, 'module'):
state_dict = net.module.state_dict()
else:
state_dict = net.state_dict()
torch.save(state_dict, './res/model_final.pth')
logger.info('\nTraining done, model saved to {}\n\n'.format(
'./res/model_final.pth'))
False)
ids.extend(id_sam.tolist())
else:
id_sam = np.random.choice(self.person_infos[pid], self.K,
True)
ids.extend(id_sam.tolist())
yield ids
def __len__(self):
return self.iter_num
if __name__ == "__main__":
from torch.utils.data import DataLoader
from market1501 import Market1501
import cv2
ds = Market1501('./dataset/Market-1501-v15.09.15/bounding_box_train')
sampler1 = BalancedSampler(ds, 2, 4)
sampler2 = BalancedSampler(ds, 2, 4)
dl1 = DataLoader(ds, batch_sampler=sampler1, num_workers=1)
dl2 = DataLoader(ds, batch_sampler=sampler2, num_workers=1)
for jj in range(2):
for i, ((imgs1, lbs1, ids1), (imgs2, lbs2,
ids2)) in enumerate(zip(dl1, dl2)):
print(i)
print(lbs1)
print(lbs2)
if i == 4: break
def run(args):
gpuId, epochs, weight_decay, batch_id, batch_image, lr_1, lr_2, erasing_p, sampling, exp_dir, trainset_name, cd_trainset_name, testset_names, rand_crop, head_1part_stride = \
args.gpuId, args.epochs, args.weight_decay, args.batch_id, args.batch_image, args.lr_1, args.lr_2, args.erasing_p, args.sampling, args.exp_dir, args.trainset_name, args.cd_trainset_name, args.testset_names, args.rand_crop, args.head_1part_stride
DEVICE = torch.device("cuda:" +
gpuId if torch.cuda.is_available() else "cpu")
print(DEVICE)
num_workers = 4
batch_test = 64 #32
train_list = [
transforms.Resize((400, 144)),
transforms.RandomCrop((384, 128))
] if rand_crop else [transforms.Resize((384, 128))]
train_list += [
transforms.ToTensor(),
]
if erasing_p > 0:
train_list = train_list + [
RandomErasing(probability=erasing_p, mean=[0.0, 0.0, 0.0])
]
train_list += [
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]
train_transform = transforms.Compose(train_list)
if trainset_name in ['market1501', 'cuhk03', 'duke']:
root = get_dataset_root(trainset_name)
train_dataset = Market1501(
root + '/bounding_box_train',
transform=train_transform,
training=True,
kpt_file=trainset_name + '-kpt.pkl' if args.pap else None,
ps_dir=root + '_ps_label' if args.src_ps_lw > 0 else None)
elif trainset_name in ['msmt17']:
train_dataset = MSMT17(transform=train_transform,
training=True,
use_kpt=args.pap,
use_ps=args.src_ps_lw > 0,
split='train')
else:
raise ValueError('Invalid train set {}'.format(trainset_name))
train_loader = DataLoader(train_dataset,
sampler=RandomIdSampler(train_dataset,
batch_image=batch_image),
batch_size=batch_id * batch_image,
num_workers=num_workers,
drop_last=True)
if args.cd_ps_lw > 0:
if cd_trainset_name in ['market1501', 'cuhk03', 'duke']:
cd_train_dataset = Market1501(get_dataset_root(cd_trainset_name) +
'/bounding_box_train',
transform=train_transform,
training=True,
kpt_file=None,
ps_dir=cd_trainset_name + '-ps')
elif cd_trainset_name in ['msmt17']:
cd_train_dataset = MSMT17(transform=train_transform,
training=True,
use_kpt=False,
use_ps=True)
else:
raise ValueError(
'Invalid cd train set {}'.format(cd_trainset_name))
cd_train_loader = InfiniteNextBatch(
DataLoader(cd_train_dataset,
batch_size=args.cd_train_batch_size,
num_workers=num_workers,
drop_last=True))
test_transform = transforms.Compose([
transforms.Resize((384, 128)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
test_flip_transform = transforms.Compose([
transforms.Resize((384, 128)), functional.hflip,
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
def make_test_loader_M_C_D(root, name):
query_dataset = Market1501(root + '/query',
transform=test_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
query_flip_dataset = Market1501(root + '/query',
transform=test_flip_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
query_loader = DataLoader(query_dataset,
batch_size=batch_test,
num_workers=num_workers)
query_flip_loader = DataLoader(query_flip_dataset,
batch_size=batch_test,
num_workers=num_workers)
test_dataset = Market1501(root + '/bounding_box_test',
transform=test_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
test_flip_dataset = Market1501(root + '/bounding_box_test',
transform=test_flip_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
test_loader = DataLoader(test_dataset,
batch_size=batch_test,
num_workers=num_workers)
test_flip_loader = DataLoader(test_flip_dataset,
batch_size=batch_test,
num_workers=num_workers)
return query_loader, query_flip_loader, test_loader, test_flip_loader
def make_test_loader_MS_PR_PI(name):
if name == 'msmt17':
dclass = MSMT17
elif name == 'partial_reid':
dclass = PartialREID
elif name == 'partial_ilids':
dclass = PartialiLIDs
else:
raise ValueError('Invalid dataset name {}'.format(name))
q_set = dclass(transform=test_transform,
training=False,
use_kpt=args.pap,
use_ps=False,
split='query')
q_flip_set = dclass(transform=test_flip_transform,
training=False,
use_kpt=args.pap,
use_ps=False,
split='query')
q_loader = DataLoader(q_set,
batch_size=batch_test,
num_workers=num_workers)
q_flip_loader = DataLoader(q_flip_set,
batch_size=batch_test,
num_workers=num_workers)
g_set = dclass(transform=test_transform,
training=False,
use_kpt=args.pap,
use_ps=False,
split='gallery')
g_flip_set = dclass(transform=test_flip_transform,
training=False,
use_kpt=args.pap,
use_ps=False,
split='gallery')
g_loader = DataLoader(g_set,
batch_size=batch_test,
num_workers=num_workers)
g_flip_loader = DataLoader(g_flip_set,
batch_size=batch_test,
num_workers=num_workers)
return q_loader, q_flip_loader, g_loader, g_flip_loader
def make_test_loader(name):
if name in ['market1501', 'cuhk03', 'duke']:
return make_test_loader_M_C_D(get_dataset_root(name), name)
elif name in ['msmt17', 'partial_reid', 'partial_ilids']:
return make_test_loader_MS_PR_PI(name)
test_loaders = [make_test_loader(name) for name in testset_names]
mgn = MGN(len(train_dataset.unique_ids), args)
if torch.cuda.device_count() > 1:
mgn = nn.DataParallel(mgn)
mgn = mgn.to(DEVICE)
vanilla_cross_entropy_loss = nn.CrossEntropyLoss()
cross_entropy_loss = nn.CrossEntropyLoss(reduce=False)
triplet_semihard_loss = TripletSemihardLoss(
margin=0.1,
DEVICE=DEVICE,
sampling=sampling,
batch_id=batch_id,
batch_image=batch_image) #batch_hard, .'curriculum'
ps_loss = PSLoss()
optimizer_start1 = optim.SGD(mgn.parameters(),
lr=lr_1,
momentum=0.9,
weight_decay=weight_decay)
optimizer_start2 = optim.SGD(mgn.parameters(),
lr=lr_2,
momentum=0.9,
weight_decay=weight_decay)
scheduler_1 = optim.lr_scheduler.MultiStepLR(optimizer_start1, [140, 180],
gamma=0.1)
scheduler_2 = optim.lr_scheduler.MultiStepLR(
optimizer_start2, [140, 180], gamma=0.1) # best [140, 180] [120, 160]
def get_model_input(inputs, target):
dic = {'im': inputs.to(DEVICE)}
if 'pap_mask_2p' in target:
dic['pap_mask_2p'] = target['pap_mask_2p'].to(DEVICE)
dic['pap_mask_3p'] = target['pap_mask_3p'].to(DEVICE)
return dic
def extract_loader_feat(loader, verbose=False):
feat = []
vis = []
i = 0
for inputs, target in loader:
if verbose:
print(i)
i += 1
with torch.no_grad():
output = mgn(get_model_input(inputs, target))
feat.append(output[1].detach().cpu().numpy())
if args.pap:
vis_ = np.concatenate([
np.ones([len(output[1]), 3]),
torch.stack(output[5 + 3 + 5:5 + 3 + 5 + 5],
1).detach().cpu().numpy()
], 1)
vis.append(vis_)
feat = np.concatenate(feat)
vis = np.concatenate(vis) if args.pap else None
return feat, vis
def test(query_loader,
query_flip_loader,
test_loader,
test_flip_loader,
trainset_name,
testset_name,
epoch,
verbose=False):
cache_file = '{}/feat_cache-{}_to_{}.pkl'.format(
exp_dir, trainset_name, testset_name)
if args.use_feat_cache:
assert os.path.exists(
cache_file), "Feature cache file {} does not exist!".format(
cache_file)
query_2, q_vis, query_flip_2, q_vis, test_2, test_vis, test_flip_2, test_vis, q_ids, q_cams, g_ids, g_cams = load_pickle(
cache_file)
else:
query_2, q_vis = extract_loader_feat(query_loader, verbose=verbose)
query_flip_2, q_vis = extract_loader_feat(query_flip_loader,
verbose=verbose)
test_2, test_vis = extract_loader_feat(test_loader,
verbose=verbose)
test_flip_2, test_vis = extract_loader_feat(test_flip_loader,
verbose=verbose)
q_ids = query_loader.dataset.ids
q_cams = query_loader.dataset.cameras
g_ids = test_loader.dataset.ids
g_cams = test_loader.dataset.cameras
save_pickle([
query_2, q_vis, query_flip_2, q_vis, test_2, test_vis,
test_flip_2, test_vis, q_ids, q_cams, g_ids, g_cams
], cache_file)
if args.test_which_feat > 0:
# TODO: implement for pap
idx = args.test_which_feat
query_2 = query_2[:, 256 * idx - 256:256 * idx]
query_flip_2 = query_flip_2[:, 256 * idx - 256:256 * idx]
test_2 = test_2[:, 256 * idx - 256:256 * idx]
test_flip_2 = test_flip_2[:, 256 * idx - 256:256 * idx]
query = normalize(query_2 + query_flip_2)
test = normalize(test_2 + test_flip_2)
if verbose:
print('query.shape:', query.shape)
print('test.shape:', test.shape)
if args.pap:
print('q_vis.shape:', q_vis.shape)
print('test_vis.shape:', test_vis.shape)
if args.pap:
dist_1 = compute_dist_with_visibility(query,
test,
q_vis,
test_vis,
dist_type='euclidean',
avg_by_vis_num=False)
else:
dist_1 = cdist(query, test)
r_1 = cmc(dist_1,
q_ids,
g_ids,
q_cams,
g_cams,
separate_camera_set=False,
single_gallery_shot=False,
first_match_break=True)
m_ap_1 = mean_ap(dist_1, q_ids, g_ids, q_cams, g_cams)
print('EPOCH [%d] %s -> %s: mAP=%f, r@1=%f, r@3=%f, r@5=%f, r@10=%f' %
(epoch + 1, trainset_name, testset_name, m_ap_1, r_1[0], r_1[2],
r_1[4], r_1[9]))
if args.only_test:
mgn.eval()
if not args.use_feat_cache:
if args.model_weight_file:
model_weight_file = args.model_weight_file
else:
model_weight_file = '{}/model_weight.pth'.format(exp_dir)
load_model_weight((mgn.module if hasattr(mgn, 'module') else mgn),
model_weight_file)
for name, test_loader in zip(testset_names, test_loaders):
test(test_loader[0],
test_loader[1],
test_loader[2],
test_loader[3],
trainset_name,
name,
-1,
verbose=False)
exit()
for epoch in range(epochs):
mgn.train()
scheduler_1.step()
scheduler_2.step()
running_loss = 0.0
running_loss_1 = 0.0
running_loss_2 = 0.0
if epoch < 20:
optimizer_1 = optim.SGD(mgn.parameters(),
lr=0.01 + 0.0045 * epoch,
momentum=0.9,
weight_decay=weight_decay)
optimizer_2 = optim.SGD(mgn.parameters(),
lr=0.001 + 0.00045 * epoch,
momentum=0.9,
weight_decay=weight_decay)
else:
optimizer_1 = optimizer_start1
optimizer_2 = optimizer_start2
for i, data in enumerate(train_loader):
inputs, target = data
inputs = inputs.to(DEVICE)
for k, v in target.items():
target[k] = v.to(DEVICE)
labels = target['id']
outputs = mgn(get_model_input(inputs, target))
optimizer_1.zero_grad()
if args.pap:
losses_1 = [
vanilla_cross_entropy_loss(output, labels)
for output in outputs[5:5 + 3]
] + [(cross_entropy_loss(output, labels) * v).sum() /
(v.sum() + 1e-12)
for output, v in zip(outputs[5 + 3:5 + 3 +
5], outputs[5 + 3 + 5:5 + 3 +
5 + 5])]
else:
losses_1 = [
vanilla_cross_entropy_loss(output, labels)
for output in outputs[5:5 + 8]
]
loss_1 = sum(losses_1) / len(losses_1)
psl = 0
if args.src_ps_lw > 0:
psl = (ps_loss(outputs[-3], target['ps_label']) +
ps_loss(outputs[-2], target['ps_label']) +
ps_loss(outputs[-1], target['ps_label'])) / 3.
(loss_1 + psl * args.src_ps_lw).backward()
if args.cd_ps_lw > 0:
cd_inputs, cd_targets = cd_train_loader.next_batch()
cd_inputs = cd_inputs.to(DEVICE)
for k, v in cd_targets.items():
cd_targets[k] = v.to(DEVICE)
pap_old = args.pap
args.pap = False
outputs = mgn(get_model_input(cd_inputs, cd_targets))
args.pap = pap_old
cd_psl = (ps_loss(outputs[-3], cd_targets['ps_label']) +
ps_loss(outputs[-2], cd_targets['ps_label']) +
ps_loss(outputs[-1], cd_targets['ps_label'])) / 3.
(cd_psl * args.cd_ps_lw).backward()
optimizer_1.step()
outputs = mgn(get_model_input(inputs, target))
optimizer_2.zero_grad()
losses_2 = [
triplet_semihard_loss(output, labels, epoch)
for output in outputs[2:5]
]
loss_2 = sum(losses_2) / len(losses_2)
psl = 0
if args.src_ps_lw > 0:
psl = (ps_loss(outputs[-3], target['ps_label']) +
ps_loss(outputs[-2], target['ps_label']) +
ps_loss(outputs[-1], target['ps_label'])) / 3.
(loss_2 + psl * args.src_ps_lw).backward()
if args.cd_ps_lw > 0:
cd_inputs, cd_targets = cd_train_loader.next_batch()
cd_inputs = cd_inputs.to(DEVICE)
for k, v in cd_targets.items():
cd_targets[k] = v.to(DEVICE)
pap_old = args.pap
args.pap = False
outputs = mgn(get_model_input(cd_inputs, cd_targets))
args.pap = pap_old
cd_psl = (ps_loss(outputs[-3], cd_targets['ps_label']) +
ps_loss(outputs[-2], cd_targets['ps_label']) +
ps_loss(outputs[-1], cd_targets['ps_label'])) / 3.
(cd_psl * args.cd_ps_lw).backward()
optimizer_2.step()
running_loss_1 += loss_1.item()
running_loss_2 += loss_2.item()
running_loss = running_loss + (loss_1.item() + loss_2.item()) / 2.0
print('%d/%d - %d/%d - loss: %f - ps_loss: %f - cd_ps_loss: %f' %
(epoch + 1, epochs, i, len(train_loader),
(loss_1.item() + loss_2.item()) / 2,
psl.item() if isinstance(psl, torch.Tensor) else 0,
cd_psl.item() if args.cd_ps_lw > 0 else 0))
print('epoch: %d/%d - loss1: %f' %
(epoch + 1, epochs, running_loss_1 / len(train_loader)))
print('epoch: %d/%d - loss2: %f' %
(epoch + 1, epochs, running_loss_2 / len(train_loader)))
# if (epoch + 1) % 50 == 0:
# model_weight_file = '{}/model_weight.pth'.format(exp_dir)
# save_model(mgn, model_weight_file)
# mgn.eval()
# for name, test_loader in zip(testset_names, test_loaders):
# test(test_loader[0], test_loader[1], test_loader[2], test_loader[3], trainset_name, name, epoch)
model_weight_file = '{}/model_weight.pth'.format(exp_dir)
save_model(mgn, model_weight_file)
mgn.eval()
for name, test_loader in zip(testset_names, test_loaders):
test(test_loader[0], test_loader[1], test_loader[2], test_loader[3],
trainset_name, name, epoch)
def embed():
## load checkpoint
res_pth = './res'
mod_pth = osp.join(res_pth, 'model_final.pkl')
net = Network_D()
net.load_state_dict(torch.load(mod_pth))
net.cuda()
net.eval()
## data loader
query_set = Market1501('./dataset/Market-1501-v15.09.15/query',
is_train=False)
gallery_set = Market1501(
'./dataset/Market-1501-v15.09.15/bounding_box_test', is_train=False)
query_loader = DataLoader(query_set,
batch_size=32,
num_workers=4,
drop_last=False)
gallery_loader = DataLoader(gallery_set,
batch_size=32,
num_workers=4,
drop_last=False)
## embed
logger.info('embedding query set ...')
query_pids = []
query_camids = []
query_embds = []
for i, (im, _, ids) in enumerate(tqdm(query_loader)):
embds = []
for crop in im:
crop = crop.cuda()
embds.append(net(crop).detach().cpu().numpy())
embed = sum(embds) / len(embds)
pid = ids[0].numpy()
camid = ids[1].numpy()
query_embds.append(embed)
query_pids.extend(pid)
query_camids.extend(camid)
query_embds = np.vstack(query_embds)
query_pids = np.array(query_pids)
query_camids = np.array(query_camids)
logger.info('embedding gallery set ...')
gallery_pids = []
gallery_camids = []
gallery_embds = []
for i, (im, _, ids) in enumerate(tqdm(gallery_loader)):
embds = []
for crop in im:
crop = crop.cuda()
embds.append(net(crop).detach().cpu().numpy())
embed = sum(embds) / len(embds)
pid = ids[0].numpy()
camid = ids[1].numpy()
gallery_embds.append(embed)
gallery_pids.extend(pid)
gallery_camids.extend(camid)
gallery_embds = np.vstack(gallery_embds)
gallery_pids = np.array(gallery_pids)
gallery_camids = np.array(gallery_camids)
## dump embeds results
embd_res = (query_embds, query_pids, query_camids, gallery_embds,
gallery_pids, gallery_camids)
with open('./res/embds.pkl', 'wb') as fw:
pickle.dump(embd_res, fw)
logger.info('embedding done, dump to: ./res/embds.pkl')
return embd_res
def run():
batch_size = 32
train_transform = transforms.Compose([
transforms.Resize(144, interpolation=3),
transforms.RandomCrop((256, 128)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
test_transform = transforms.Compose([
transforms.Resize((288, 144), interpolation=3),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
test_flip_transform = transforms.Compose([
transforms.Resize((288, 144), interpolation=3),
functional.hflip,
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
train_dataset = Market1501(root + '/bounding_box_train',
transform=train_transform)
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True)
query_dataset = Market1501(root + '/query', transform=test_transform)
query_flip_dataset = Market1501(root + '/query',
transform=test_flip_transform)
query_loader = DataLoader(query_dataset,
batch_size=batch_size,
shuffle=False)
query_flip_loader = DataLoader(query_flip_dataset,
batch_size=batch_size,
shuffle=False)
test_dataset = Market1501(root + '/bounding_box_test',
transform=test_transform)
test_flip_dataset = Market1501(root + '/bounding_box_test',
transform=test_flip_transform)
test_loader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False)
test_flip_loader = DataLoader(test_flip_dataset,
batch_size=batch_size,
shuffle=False)
ide = IDE(num_classes=len(train_dataset.unique_ids)).to(DEVICE)
criterion = nn.CrossEntropyLoss()
params = [
{
'params': ide.backbone.parameters(),
'lr': 0.01
},
{
'params': ide.classifier.parameters(),
'lr': 0.1
},
]
optimizer = optim.SGD(params,
momentum=0.9,
weight_decay=5e-4,
nesterov=True)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
epochs = 50
for epoch in range(epochs):
ide.train()
scheduler.step()
running_loss = 0.0
for i, data in enumerate(train_loader):
inputs, labels = data
inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
optimizer.zero_grad()
outputs = ide(inputs)
loss = criterion(outputs[1], labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print('%d/%d - %d/%d - loss: %f' %
(epoch, epochs, i, len(train_loader), loss.item()))
print('epoch: %d/%d - loss: %f' %
(epoch, epochs, running_loss / len(train_loader)))
if epoch % 10 == 9:
ide.eval()
query = np.concatenate([
ide(inputs.to(DEVICE))[0].detach().cpu().numpy()
for inputs, _ in query_loader
])
query_flip = np.concatenate([
ide(inputs.to(DEVICE))[0].detach().cpu().numpy()
for inputs, _ in query_flip_loader
])
test = np.concatenate([
ide(inputs.to(DEVICE))[0].detach().cpu().numpy()
for inputs, _ in test_loader
])
test_flip = np.concatenate([
ide(inputs.to(DEVICE))[0].detach().cpu().numpy()
for inputs, _ in test_flip_loader
])
# dist = cdist((query + query_flip) / 2., (test + test_flip) / 2.)
dist = cdist(normalize(query + query_flip),
normalize(test + test_flip))
r = cmc(dist,
query_dataset.ids,
test_dataset.ids,
query_dataset.cameras,
test_dataset.cameras,
separate_camera_set=False,
single_gallery_shot=False,
first_match_break=True)
m_ap = mean_ap(dist, query_dataset.ids, test_dataset.ids,
query_dataset.cameras, test_dataset.cameras)
print('epoch[%d]: mAP=%f, r@1=%f, r@3=%f, r@5=%f, r@10=%f' %
(epoch + 1, m_ap, r[0], r[2], r[4], r[9]))
