def get_data(name, data_dir, num_bn_sample):
dataset = None
# Datasets
if name == 'market1501':
dataset_name = 'market1501'
dataset = data_manager.init_imgreid_dataset(root=data_dir,
name=dataset_name)
dataset.images_dir = ''
elif name == 'market_sct' or name == 'market_sct_tran' or name == 'duke_sct' or name == 'duke_sct_tran':
dataset_name = name
dataset = data_manager.init_imgreid_dataset(
root=data_dir, name=dataset_name, num_bn_sample=num_bn_sample)
dataset.images_dir = ''
assert dataset is not None, "dataset is none"
return dataset
def main():
global args, best_rank1
torch.set_num_threads(args.workers)
torch.manual_seed(args.seed)
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if args.save_dir:
save_dir = args.save_dir
else:
save_dir = osp.join('logs', 'dukemtmcreid_s1')
writer = SummaryWriter(log_dir=save_dir)
sys.stdout = Logger(osp.join(save_dir, 'log_train.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_imgreid_dataset(root=args.root,
name=args.dataset)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
# decompose tracklets into images for image-based training
train_data = dataset.get_train_tracklets()
new_train = []
for img_paths, pid, camid, _, _, _ in train_data:
for img_path in img_paths:
new_train.append((img_path, pid, camid))
trainloader = DataLoader(
ImageDataset(new_train, transform=transform_train),
sampler=ClassIdentitySampler(new_train, args.train_batch,
args.num_instances),
batch_size=args.train_batch,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transform_test),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, loss={'htri'})
print("Model size: {:.3f} M".format(count_num_param(model)))
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
if use_gpu:
model = nn.DataParallel(model).cuda()
start_time = time.time()
train_time = 0
best_epoch = args.start_epoch
print("==> Start training")
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
# test before train
if epoch == 0:
mAP, rank1, rank5, rank10, rank20 = test(model, queryloader,
galleryloader, use_gpu)
res_dict = {
'mAP': mAP,
'rank-1': rank1,
'rank-5': rank5,
'rank-10': rank10,
'rank-20': rank20,
}
writer.add_scalars('scalar/precision', res_dict, epoch)
train(epoch, model, criterion_htri, optimizer, trainloader, use_gpu,
writer)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (
epoch + 1) % args.eval_step == 0 or (epoch +
1) == args.max_epoch:
print("==> Test")
mAP, rank1, rank5, rank10, rank20 = test(model, queryloader,
galleryloader, use_gpu)
res_dict = {
'mAP': mAP,
'rank-1': rank1,
'rank-5': rank5,
'rank-10': rank10,
'rank-20': rank20,
}
writer.add_scalars('scalar/precision', res_dict, epoch + 1)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
if is_best:
save_checkpoint(
{
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
},
fpath=osp.join(save_dir, 's1_best_model' + '.pth.tar'))
train_data = dataset.get_train_tracklets()
new_train = []
for img_paths, pid, camid, _, _, _ in train_data:
for img_path in img_paths:
new_train.append((img_path, pid, camid))
trainloader = DataLoader(
ImageDataset(new_train, transform=transform_train),
sampler=ClassIdentitySampler(new_train, args.train_batch,
args.num_instances),
batch_size=args.train_batch,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
save_checkpoint(
{
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, False, osp.join(save_dir, 's1_checkpoint_final' + '.pth.tar'))
print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(
best_rank1, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".
format(elapsed, train_time))
def get_data(name,
split_id,
data_dir,
height,
width,
batch_size,
num_instances,
workers,
combine_trainval,
eval_rerank=False):
## Datasets
if name == 'cuhk03labeled':
dataset_name = 'cuhk03'
dataset = data_manager.init_imgreid_dataset(
root=data_dir,
name=dataset_name,
split_id=split_id,
cuhk03_labeled=True,
cuhk03_classic_split=False,
)
dataset.images_dir = osp.join(data_dir, '/CUHK03_New/images_labeled/')
elif name == 'cuhk03detected':
dataset_name = 'cuhk03'
dataset = data_manager.init_imgreid_dataset(
root=data_dir,
name=dataset_name,
split_id=split_id,
cuhk03_labeled=False,
cuhk03_classic_split=False,
)
dataset.images_dir = osp.join(data_dir, '/CUHK03_New/images_detected/')
## Num. of training IDs
num_classes = dataset.num_train_pids
train_transformer = T.Compose([
T.Random2DTranslation(height, width),
T.RandomHorizontalFlip(),
T.ToTensor(),
])
test_transformer = T.Compose([
T.RectScale(height, width),
T.ToTensor(),
])
train_loader = DataLoader(Preprocessor(dataset.train,
root=dataset.images_dir,
transform=train_transformer),
batch_size=batch_size,
num_workers=workers,
sampler=RandomIdentitySampler(
dataset.train, num_instances),
pin_memory=True,
drop_last=True)
query_loader = DataLoader(Preprocessor(dataset.query,
root=dataset.images_dir,
transform=test_transformer),
batch_size=batch_size,
num_workers=workers,
shuffle=False,
pin_memory=True)
gallery_loader = DataLoader(Preprocessor(dataset.gallery,
root=dataset.images_dir,
transform=test_transformer),
batch_size=batch_size,
num_workers=workers,
shuffle=False,
pin_memory=True)
return dataset, num_classes, train_loader, query_loader, gallery_loader
def main():
global args, best_rank1
torch.manual_seed(args.seed)
if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if args.save_dir:
save_dir = args.save_dir
else:
save_dir = osp.join('logs', 'dukemtmcreid_s2')
if not args.evaluate:
writer = SummaryWriter(log_dir=save_dir)
sys.stdout = Logger(osp.join(save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_imgreid_dataset(root=args.root, name=args.dataset)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
print("Initializing model: {}".format(args.arch))
model_s1 = models.init_model(name=args.arch, loss={'htri'})
if args.scratch:
model_0 = models.init_model(name=args.arch, loss={'htri'})
print("Model size: {:.3f} M".format(count_num_param(model_s1)))
if args.resume and check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model_s1.load_state_dict(checkpoint['state_dict'])
best_rank1 = checkpoint['rank1']
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, best_rank1))
if args.load_weights and check_isfile(args.load_weights):
# load pretrained weights but ignore layers that don't match in size
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model_s1.state_dict()
pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()}
model_dict.update(pretrain_dict)
model_s1.load_state_dict(model_dict)
print("Loaded pretrained weights from '{}'".format(args.load_weights))
if use_gpu:
model_s1 = nn.DataParallel(model_s1).cuda()
if args.scratch:
model_0 = nn.DataParallel(model_0).cuda()
if args.scratch:
optimizer = init_optim(args.optim, model_0.parameters(), args.lr, args.weight_decay)
else:
optimizer = init_optim(args.optim, model_s1.parameters(), args.lr, args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma)
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
criterion_htri = TripletLoss(margin=args.margin)
if args.evaluate:
print("Evaluate only")
mAP, rank1, rank5, rank10, rank20 = test(model, queryloader, galleryloader, use_gpu)
return
start_time = time.time()
train_time = 0
best_epoch = args.start_epoch
print("==> Start training")
train_pair_cnt = []
train_pair_cnt_true = []
for epoch in range(args.start_epoch, args.max_epoch):
if epoch == 0:
if args.scratch:
model = model_0
else:
model = model_s1
# association
if epoch%args.update_step == 0:
print("extract tracklets features ...")
feats_cams, tids_cams, camids_cams, tmins_cams, tmaxs_cams, otids_cams = \
[], [], [], [], [], []
for cam_tracklets in dataset.tracklets:
# print(cam_tracklets[0])
# break
trackletsloader = DataLoader(
VideoDataset_time(cam_tracklets, seq_len=args.seq_len, sample='evenly', transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
if epoch == 0:
feats, tids, camids, tmins, tmaxs, otids = \
extract_feat(model_s1, trackletsloader, args.pool, use_gpu)
else:
feats, tids, camids, tmins, tmaxs, otids = \
extract_feat(model, trackletsloader, args.pool, use_gpu)
feats_cams.append(feats)
tids_cams.append(tids)
camids_cams.append(camids)
tmins_cams.append(tmins)
tmaxs_cams.append(tmaxs)
otids_cams.append(otids)
cnt, cnt_true = dataset.update_train(epoch, feats_cams, camids_cams, tmins_cams, tmaxs_cams, save_dir, k=args.kmeans, nosp=args.nosp, sigma=args.sigma)
train_pair_cnt.append(cnt)
train_pair_cnt_true.append(cnt_true)
new_train = []
train_data_s2 = dataset.get_train_tracklet_s2()
for img_paths, pid, gid in train_data_s2:
for img_path in img_paths:
new_train.append((img_path, pid, gid))
trainloader = DataLoader(
ImageDataset(new_train, transform=transform_train),
sampler=ClassIdentitySampler(new_train, args.train_batch, args.num_instances),
batch_size=args.train_batch, num_workers=args.workers,
pin_memory=pin_memory, drop_last=True,
)
start_train_time = time.time()
# test before train
if epoch == 0:
mAP, rank1, rank5, rank10, rank20 = test(model, queryloader, galleryloader, use_gpu)
res_dict = {
'mAP': mAP,
'rank-1': rank1,
'rank-5': rank5,
'rank-10': rank10,
'rank-20': rank20,
}
writer.add_scalars('scalar/precision', res_dict, epoch)
train(epoch, model, criterion_htri, optimizer, trainloader, use_gpu, writer)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch:
print("==> Test")
mAP, rank1, rank5, rank10, rank20 = test(model, queryloader, galleryloader, use_gpu)
res_dict = {
'mAP': mAP,
'rank-1': rank1,
'rank-5': rank5,
'rank-10': rank10,
'rank-20': rank20,
}
writer.add_scalars('scalar/precision', res_dict, epoch+1)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
if is_best:
save_checkpoint({
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, fpath=osp.join(save_dir, 's2_best_model' + '.pth.tar'))
save_checkpoint({
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, False, osp.join(save_dir, 's2_checkpoint_final' + '.pth.tar'))
print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(best_rank1, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time))
train_pair_cnt = np.array(train_pair_cnt)
train_pair_cnt_true = np.array(train_pair_cnt_true)
assert train_pair_cnt.shape == train_pair_cnt_true.shape
for i in range(train_pair_cnt.shape[0]):
print('iter %d'%(i))
print(train_pair_cnt[i], train_pair_cnt[i].sum())
print(train_pair_cnt_true[i], train_pair_cnt_true[i].sum(), train_pair_cnt_true[i].sum()/train_pair_cnt[i].sum())
def get_data(name, data_dir, height, width, batch_size, num_bn_sample,
num_instances, workers):
# Datasets
if name == 'market1501':
dataset_name = 'market1501'
dataset = data_manager.init_imgreid_dataset(root=data_dir,
name=dataset_name)
dataset.images_dir = ''
train_transformer = T.Compose([
T.Random2DTranslation(height, width),
T.RandomHorizontalFlip(),
T.ToTensor(),
])
test_transformer = T.Compose([
T.RectScale(height, width),
T.ToTensor(),
])
train_loader = DataLoader(Preprocessor(dataset.train,
root=dataset.images_dir,
transform=train_transformer),
batch_size=batch_size,
num_workers=workers,
sampler=RandomIdentitySampler(
dataset.train, num_instances),
pin_memory=True,
drop_last=True)
elif name == 'market_sct' or name == 'market_sct_tran' or name == 'duke_sct' or name == 'duke_sct_tran':
dataset_name = name
dataset = data_manager.init_imgreid_dataset(
root=data_dir, name=dataset_name, num_bn_sample=num_bn_sample)
dataset.images_dir = ''
# dataset_name = name
pin_memory = True
collateFn = NormalCollateFn()
train_loader = DataLoader(
data_manager.init_datafolder(
dataset_name,
dataset.train,
TrainTransform(height, width),
),
batch_sampler=RandomIdentityCameraSampler(dataset.train,
num_instances,
batch_size),
num_workers=workers,
pin_memory=pin_memory,
collate_fn=collateFn,
)
# train_loader = DataLoader(
# data_manager.init_datafolder(dataset_name, dataset.train,
# TrainTransform(height, width),
# ),
# batch_sampler=RandomIdentityUniqueCameraSampler(dataset.train, num_instances, batch_size),
# num_workers=workers,
# pin_memory=pin_memory, collate_fn=collateFn,
# )
# Num. of training IDs
num_classes = dataset.num_train_pids
return dataset, num_classes, train_loader