def get_train_loader(dataset, height, width, batch_size, workers,
num_instances, iters, trainset=None):
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.RandomHorizontalFlip(p=0.5),
T.Pad(10),
T.RandomCrop((height, width)),
T.ToTensor(),
normalizer,
T.RandomErasing(probability=0.5, mean=[0.485, 0.456, 0.406])
])
train_set = dataset.train if trainset is None else trainset
rmgs_flag = num_instances > 0
if rmgs_flag:
sampler = RandomMultipleGallerySampler(train_set, num_instances)
else:
sampler = None
train_loader = IterLoader(
DataLoader(Preprocessor(train_set, root=dataset.images_dir,
transform=train_transformer, mutual=False),
batch_size=batch_size, num_workers=workers, sampler=sampler,
shuffle=not rmgs_flag, pin_memory=True, drop_last=True), length=iters)
return train_loader
def get_data(name, data_dir, height, width, batch_size, workers, num_instances, iters=200):
root = osp.join(data_dir, name)
dataset = datasets.create(name, root)
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_set = dataset.train
num_classes = dataset.num_train_pids
train_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.RandomHorizontalFlip(p=0.5),
T.Pad(10),
T.RandomCrop((height, width)),
T.ToTensor(),
normalizer
])
test_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.ToTensor(),
normalizer
])
rmgs_flag = num_instances > 0
if rmgs_flag:
sampler = RandomMultipleGallerySampler(train_set, num_instances)
else:
sampler = None
train_loader = IterLoader(
DataLoader(Preprocessor(train_set, root=dataset.images_dir,
transform=train_transformer),
batch_size=batch_size, num_workers=workers, sampler=sampler,
shuffle=not rmgs_flag, pin_memory=True, drop_last=True), length=iters)
test_loader = DataLoader(
Preprocessor(list(set(dataset.query) | set(dataset.gallery)),
root=dataset.images_dir, transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
return dataset, num_classes, train_loader, test_loader
def get_data(data_dir, source, target, height, width, batch_size, num_instance=2, workers=8):
dataset = DA(data_dir, source, target)
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
num_classes = dataset.num_train_ids
train_transformer = T.Compose([
T.Resize((256, 128), interpolation=3),
T.Pad(10),
T.RandomCrop((256,128)),
T.RandomHorizontalFlip(0.5),
T.RandomRotation(5),
T.ToTensor(),
normalizer,
])
test_transformer = T.Compose([
T.Resize((256, 128), interpolation=3),
T.ToTensor(),
normalizer,
])
source_train_loader = DataLoader(
Preprocessor_occluded(dataset.source_train, root=osp.join(dataset.source_images_dir, dataset.source_train_path),
transform=train_transformer, train=True),
batch_size=batch_size, num_workers=workers,
sampler=IdentitySampler(dataset.source_train, num_instance),
pin_memory=True, drop_last=True)
query_loader = DataLoader(
Preprocessor_occluded(dataset.query,
root=osp.join(dataset.target_images_dir, dataset.query_path), transform=test_transformer),
batch_size=42, num_workers=workers,
shuffle=False, pin_memory=True)
gallery_loader = DataLoader(
Preprocessor_occluded(dataset.gallery,
root=osp.join(dataset.target_images_dir, dataset.gallery_path), transform=test_transformer),
batch_size=42, num_workers=workers,
shuffle=False, pin_memory=True)
return dataset, num_classes, source_train_loader, query_loader, gallery_loader
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
print(args)
shutil.copy(sys.argv[0], osp.join(args.logs_dir,
osp.basename(sys.argv[0])))
# Create data loaders
if args.height is None or args.width is None:
args.height, args.width = (256, 128)
dataset, num_classes, train_loader, val_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.height,
args.width, args.batch_size * 8, args.workers,
)
# Create model
model = models.create("ft_net_inter",
num_classes=num_classes,
stride=args.stride)
# Load from checkpoint
start_epoch = 0
best_top1 = 0
top1 = 0
is_best = False
if args.checkpoint is not None:
if args.evaluate:
checkpoint = load_checkpoint(args.checkpoint)
param_dict = model.state_dict()
for k, v in checkpoint['state_dict'].items():
if 'model' in k:
param_dict[k] = v
model.load_state_dict(param_dict)
else:
model.model.load_param(args.checkpoint)
model = model.cuda()
# Distance metric
metric = None
# Evaluator
evaluator = Evaluator(model, use_cpu=args.use_cpu)
if args.evaluate:
print("Test:")
evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
return
train_transformer = [
T.Resize((args.height, args.width), interpolation=3),
T.RandomHorizontalFlip(),
T.Pad(10),
T.RandomCrop((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
T.RandomErasing(probability=0.5),
]
train_transformer = T.Compose(train_transformer)
for cluster_epoch in range(args.cluster_epochs):
# -------------------------Stage 1 intra camera training--------------------------
# Cluster and generate new dataset and model
cluster_result = get_intra_cam_cluster_result(model, train_loader,
args.class_number_stage1,
args.linkage)
cluster_datasets = [
datasets.create("cluster", osp.join(args.data_dir, args.dataset),
cluster_result[cam_id], cam_id)
for cam_id in cluster_result.keys()
]
cluster_dataloaders = [
DataLoader(Preprocessor(dataset.train_set,
root=dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size,
num_workers=args.workers,
shuffle=True,
pin_memory=False,
drop_last=True) for dataset in cluster_datasets
]
param_dict = model.model.state_dict()
model = models.create("ft_net_intra",
num_classes=[
args.class_number_stage1
for cam_id in cluster_result.keys()
],
stride=args.stride)
model_param_dict = model.model.state_dict()
for k, v in model_param_dict.items():
if k in param_dict.keys():
model_param_dict[k] = param_dict[k]
model.model.load_state_dict(model_param_dict)
model = model.cuda()
criterion = nn.CrossEntropyLoss().cuda()
# Optimizer
param_groups = make_params(model, args.lr, args.weight_decay)
optimizer = torch.optim.SGD(param_groups, momentum=0.9)
# Trainer
trainer = IntraCameraTrainer(model,
criterion,
warm_up_epoch=args.warm_up)
print("start training")
# Start training
for epoch in range(0, args.epochs_stage1):
trainer.train(
cluster_epoch,
epoch,
cluster_dataloaders,
optimizer,
print_freq=args.print_freq,
)
# -------------------------------------------Stage 2 inter camera training-----------------------------------
mix_rate = get_mix_rate(args.mix_rate,
cluster_epoch,
args.cluster_epochs,
power=args.decay_factor)
cluster_result = get_inter_cam_cluster_result(model,
train_loader,
args.class_number_stage2,
args.linkage,
mix_rate,
use_cpu=args.use_cpu)
cluster_dataset = datasets.create(
"cluster", osp.join(args.data_dir, args.dataset), cluster_result,
0)
cluster_dataloaders = DataLoader(
Preprocessor(cluster_dataset.train_set,
root=cluster_dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size_stage2,
num_workers=args.workers,
sampler=RandomIdentitySampler(cluster_dataset.train_set,
args.batch_size_stage2,
args.instances),
pin_memory=False,
drop_last=True)
param_dict = model.model.state_dict()
model = models.create("ft_net_inter",
num_classes=args.class_number_stage2,
stride=args.stride)
model.model.load_state_dict(param_dict)
model = model.cuda()
# Criterion
criterion_entropy = nn.CrossEntropyLoss().cuda()
criterion_triple = TripletLoss(margin=args.margin).cuda()
# Optimizer
param_groups = make_params(model,
args.lr * args.batch_size_stage2 / 32,
args.weight_decay)
optimizer = torch.optim.SGD(param_groups, momentum=0.9)
# Trainer
trainer = InterCameraTrainer(
model,
criterion_entropy,
criterion_triple,
warm_up_epoch=args.warm_up,
)
print("start training")
# Start training
for epoch in range(0, args.epochs_stage2):
trainer.train(cluster_epoch,
epoch,
cluster_dataloaders,
optimizer,
print_freq=args.print_freq)
if (cluster_epoch + 1) % 5 == 0:
evaluator = Evaluator(model, use_cpu=args.use_cpu)
top1, mAP = evaluator.evaluate(test_loader,
dataset.query,
dataset.gallery,
metric,
return_mAP=True)
is_best = top1 > best_top1
best_top1 = max(top1, best_top1)
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': cluster_epoch + 1,
'best_top1': best_top1,
'cluster_epoch': cluster_epoch + 1,
},
is_best,
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
if cluster_epoch == (args.cluster_epochs - 1):
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': cluster_epoch + 1,
'best_top1': best_top1,
'cluster_epoch': cluster_epoch + 1,
},
False,
fpath=osp.join(args.logs_dir, 'latest.pth.tar'))
print('\n * cluster_epoch: {:3d} top1: {:5.1%} best: {:5.1%}{}\n'.
format(cluster_epoch, top1, best_top1, ' *' if is_best else ''))
# Final test
print('Test with best model:')
checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
model.load_state_dict(checkpoint['state_dict'])
best_rank1, mAP = evaluator.evaluate(test_loader,
dataset.query,
dataset.gallery,
metric,
return_mAP=True)
def get_data(name, split_id, data_dir, height, width, batch_size,
num_instances, workers, combine_trainval, flip_prob, padding,
re_prob):
root = osp.join(data_dir, name)
print(root)
dataset = datasets.create(name, root, split_id=split_id)
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
trainvallabel = dataset.trainvallabel
train_set = dataset.trainval if combine_trainval else dataset.train
num_classes = (dataset.num_trainval_ids
if combine_trainval else dataset.num_train_ids)
train_transformer = T.Compose([
T.Resize((height, width)),
T.RandomHorizontalFlip(),
T.Pad(padding),
T.RandomCrop((height, width)),
T.ToTensor(), normalizer,
RandomErasing(probability=re_prob, mean=[0.485, 0.456, 0.406])
])
# train_transformer = T.Compose([
# T.RandomSizedRectCrop(height, width),
# T.RandomHorizontalFlip(),
# T.ToTensor(),
# normalizer,
# ])
test_transformer = T.Compose([
T.Resize((height, width)),
T.ToTensor(),
normalizer,
])
val_loader = DataLoader(Preprocessor(dataset.val,
root=dataset.images_dir,
transform=test_transformer),
batch_size=32,
num_workers=workers,
shuffle=False,
pin_memory=True)
query_loader = DataLoader(Preprocessor(list(set(dataset.query)),
root=dataset.images_dir,
transform=test_transformer),
batch_size=32,
num_workers=workers,
sampler=CamSampler(list(set(dataset.query)),
[2, 5]),
shuffle=False,
pin_memory=True)
gallery_loader = DataLoader(Preprocessor(list(set(dataset.gallery)),
root=dataset.images_dir,
transform=test_transformer),
batch_size=32,
num_workers=workers,
sampler=CamSampler(list(set(dataset.gallery)),
[0, 1, 3, 4], 4),
shuffle=False,
pin_memory=True)
train_loader = DataLoader(Preprocessor(train_set,
root=dataset.images_dir,
transform=train_transformer),
batch_size=batch_size,
num_workers=workers,
sampler=RandomIdentitySampler(
train_set, num_instances),
pin_memory=True,
drop_last=True)
return dataset, num_classes, train_loader, trainvallabel, val_loader, query_loader, gallery_loader
def get_data(data_dir, height, width, batch_size, num_instances, re=0, workers=8):
dataset = DA(data_dir)
test_dataset = TotalData(data_dir)
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
num_classes = dataset.num_source_ids
train_transformer = T.Compose([
T.Resize((256, 128), interpolation=3),
T.Pad(10),
T.RandomCrop((256,128)),
T.RandomHorizontalFlip(0.5),
T.RandomRotation(5),
T.ColorJitter(brightness=(0.5, 2.0), saturation=(0.5, 2.0), hue=(-0.1, 0.1)),
T.ToTensor(),
normalizer,
# T.RandomErasing(EPSILON=re),
])
test_transformer = T.Compose([
T.Resize((256, 128), interpolation=3),
T.ToTensor(),
normalizer,
])
# Train
source_train_loader = DataLoader(
Preprocessor(dataset.source_train,
transform=train_transformer),
batch_size=batch_size, num_workers=workers,
# shuffle=True, pin_memory=True, drop_last=True)
sampler=RandomIdentitySampler(dataset.source_train, batch_size, num_instances),
pin_memory=True, drop_last=True)
# Test
grid_query_loader = DataLoader(
Preprocessor(test_dataset.grid_query,
root=osp.join(test_dataset.grid_images_dir, test_dataset.query_path), transform=test_transformer),
batch_size=64, num_workers=4,
shuffle=False, pin_memory=True)
grid_gallery_loader = DataLoader(
Preprocessor(test_dataset.grid_gallery,
root=osp.join(test_dataset.grid_images_dir, test_dataset.gallery_path), transform=test_transformer),
batch_size=64, num_workers=4,
shuffle=False, pin_memory=True)
prid_query_loader = DataLoader(
Preprocessor(test_dataset.prid_query,
root=osp.join(test_dataset.prid_images_dir, test_dataset.query_path), transform=test_transformer),
batch_size=64, num_workers=4,
shuffle=False, pin_memory=True)
prid_gallery_loader = DataLoader(
Preprocessor(test_dataset.prid_gallery,
root=osp.join(test_dataset.prid_images_dir, test_dataset.gallery_path), transform=test_transformer),
batch_size=64, num_workers=4,
shuffle=False, pin_memory=True)
viper_query_loader = DataLoader(
Preprocessor(test_dataset.viper_query,
root=osp.join(test_dataset.viper_images_dir, test_dataset.query_path), transform=test_transformer),
batch_size=64, num_workers=4,
shuffle=False, pin_memory=True)
viper_gallery_loader = DataLoader(
Preprocessor(test_dataset.viper_gallery,
root=osp.join(test_dataset.viper_images_dir, test_dataset.gallery_path), transform=test_transformer),
batch_size=64, num_workers=4,
shuffle=False, pin_memory=True)
ilid_query_loader = DataLoader(
Preprocessor(test_dataset.ilid_query,
root=osp.join(test_dataset.ilid_images_dir, "images"), transform=test_transformer),
batch_size=64, num_workers=4,
shuffle=False, pin_memory=True)
ilid_gallery_loader = DataLoader(
Preprocessor(test_dataset.ilid_gallery,
root=osp.join(test_dataset.ilid_images_dir, "images"), transform=test_transformer),
batch_size=64, num_workers=4,
shuffle=False, pin_memory=True)
return dataset, test_dataset, num_classes, source_train_loader, grid_query_loader, grid_gallery_loader,prid_query_loader, prid_gallery_loader,viper_query_loader, viper_gallery_loader, ilid_query_loader, ilid_gallery_loader
def get_data(dataname, data_dir, model, matcher, save_path, args):
root = osp.join(data_dir, dataname)
dataset = datasets.create(dataname, root, combine_all=args.combine_all)
num_classes = dataset.num_train_ids
train_transformer = T.Compose([
T.Resize((args.height, args.width),
interpolation=InterpolationMode.BICUBIC),
T.Pad(10),
T.RandomCrop((args.height, args.width)),
T.RandomHorizontalFlip(0.5),
T.RandomRotation(5),
T.ColorJitter(brightness=(0.5, 2.0),
contrast=(0.5, 2.0),
saturation=(0.5, 2.0),
hue=(-0.1, 0.1)),
T.RandomOcclusion(args.min_size, args.max_size),
T.ToTensor(),
])
test_transformer = T.Compose([
T.Resize((args.height, args.width),
interpolation=InterpolationMode.BICUBIC),
T.ToTensor(),
])
train_path = osp.join(dataset.images_dir, dataset.train_path)
train_loader = DataLoader(
Preprocessor(dataset.train,
root=train_path,
transform=train_transformer),
batch_size=args.batch_size,
num_workers=args.workers,
sampler=GraphSampler(dataset.train, train_path, test_transformer,
model, matcher, args.batch_size,
args.num_instance, args.test_gal_batch,
args.test_prob_batch, save_path, args.gs_verbose),
pin_memory=True)
query_loader = DataLoader(Preprocessor(dataset.query,
root=osp.join(
dataset.images_dir,
dataset.query_path),
transform=test_transformer),
batch_size=args.test_fea_batch,
num_workers=args.workers,
shuffle=False,
pin_memory=True)
gallery_loader = DataLoader(Preprocessor(dataset.gallery,
root=osp.join(
dataset.images_dir,
dataset.gallery_path),
transform=test_transformer),
batch_size=args.test_fea_batch,
num_workers=args.workers,
shuffle=False,
pin_memory=True)
return dataset, num_classes, train_loader, query_loader, gallery_loader
def update_train_loader(dataset,
train_samples,
updated_label,
height,
width,
batch_size,
re,
workers,
all_img_cams,
sample_position=7):
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.RandomHorizontalFlip(p=0.5),
T.Pad(10),
T.RandomCrop((height, width)),
T.ToTensor(), normalizer,
T.RandomErasing(EPSILON=re)
])
# obtain global accumulated label from pseudo label and cameras
pure_label = updated_label[updated_label >= 0]
pure_cams = all_img_cams[updated_label >= 0]
accumulate_labels = np.zeros(pure_label.shape, pure_label.dtype)
prev_id_count = 0
id_count_each_cam = []
for this_cam in np.unique(pure_cams):
percam_labels = pure_label[pure_cams == this_cam]
unique_id = np.unique(percam_labels)
id_count_each_cam.append(len(unique_id))
id_dict = {ID: i for i, ID in enumerate(unique_id.tolist())}
for i in range(len(percam_labels)):
percam_labels[i] = id_dict[percam_labels[i]]
accumulate_labels[pure_cams ==
this_cam] = percam_labels + prev_id_count
prev_id_count += len(unique_id)
print(' sum(id_count_each_cam)= {}'.format(sum(id_count_each_cam)))
new_accum_labels = -1 * np.ones(updated_label.shape, updated_label.dtype)
new_accum_labels[updated_label >= 0] = accumulate_labels
# update sample list
new_train_samples = []
for sample in train_samples:
lbl = updated_label[sample[3]]
if lbl != -1:
assert (new_accum_labels[sample[3]] >= 0)
new_sample = sample + (lbl, new_accum_labels[sample[3]])
new_train_samples.append(new_sample)
target_train_loader = DataLoader(UnsupervisedTargetPreprocessor(
new_train_samples,
root=osp.join(dataset.target_images_dir, dataset.target_train_path),
num_cam=dataset.target_num_cam,
transform=train_transformer,
has_pseudo_label=True),
batch_size=batch_size,
num_workers=workers,
pin_memory=True,
drop_last=True,
sampler=ClassUniformlySampler(
new_train_samples,
class_position=sample_position,
k=4))
return target_train_loader, len(new_train_samples)
def get_data(name, data_dir, height, width, batch_size, workers,
combine_trainval, crop, tracking_icams, fps, re=0, num_instances=0, camstyle=0, zju=0, colorjitter=0):
# if name == 'market1501':
# root = osp.join(data_dir, 'Market-1501-v15.09.15')
# elif name == 'duke_reid':
# root = osp.join(data_dir, 'DukeMTMC-reID')
# elif name == 'duke_tracking':
# root = osp.join(data_dir, 'DukeMTMC')
# else:
# root = osp.join(data_dir, name)
if name == 'duke_tracking':
if tracking_icams != 0:
tracking_icams = [tracking_icams]
else:
tracking_icams = list(range(1, 9))
dataset = datasets.create(name, data_dir, data_type='tracking_gt', iCams=tracking_icams, fps=fps,
trainval=combine_trainval)
elif name == 'aic_tracking':
dataset = datasets.create(name, data_dir, data_type='tracking_gt', fps=fps, trainval=combine_trainval)
else:
dataset = datasets.create(name, data_dir)
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
num_classes = dataset.num_train_ids
train_transformer = T.Compose([
T.ColorJitter(brightness=0.1 * colorjitter, contrast=0.1 * colorjitter, saturation=0.1 * colorjitter, hue=0),
T.Resize((height, width)),
T.RandomHorizontalFlip(),
T.Pad(10 * crop),
T.RandomCrop((height, width)),
T.ToTensor(),
normalizer,
T.RandomErasing(probability=re),
])
test_transformer = T.Compose([
T.Resize((height, width)),
# T.RectScale(height, width, interpolation=3),
T.ToTensor(),
normalizer,
])
if zju:
train_loader = DataLoader(
Preprocessor(dataset.train, root=dataset.train_path, transform=train_transformer),
batch_size=batch_size, num_workers=workers,
sampler=ZJU_RandomIdentitySampler(dataset.train, batch_size, num_instances) if num_instances else None,
shuffle=False if num_instances else True, pin_memory=True, drop_last=False if num_instances else True)
else:
train_loader = DataLoader(
Preprocessor(dataset.train, root=dataset.train_path, transform=train_transformer),
batch_size=batch_size, num_workers=workers,
sampler=RandomIdentitySampler(dataset.train, num_instances) if num_instances else None,
shuffle=False if num_instances else True, pin_memory=True, drop_last=True)
query_loader = DataLoader(
Preprocessor(dataset.query, root=dataset.query_path, transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
gallery_loader = DataLoader(
Preprocessor(dataset.gallery, root=dataset.gallery_path, transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
if camstyle <= 0:
camstyle_loader = None
else:
camstyle_loader = DataLoader(
Preprocessor(dataset.camstyle, root=dataset.camstyle_path, transform=train_transformer),
batch_size=camstyle, num_workers=workers,
shuffle=True, pin_memory=True, drop_last=True)
return dataset, num_classes, train_loader, query_loader, gallery_loader, camstyle_loader
def get_data(data_dir, source, target, height, width, batch_size, re=0, workers=8):
dataset = DA(data_dir, source, target)
dataset_2 = DA(data_dir, target, source)
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
num_classes = dataset.num_train_ids
train_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.RandomHorizontalFlip(p=0.5),
T.Pad(10),
T.RandomCrop((height, width)),
T.ToTensor(),
normalizer,
# T.RandomErasing(EPSILON=re),
T.RandomErasing(probability=0.4, mean=[0.485, 0.456, 0.406])
])
train_transformer_2 = T.Compose([
T.Resize((height, width), interpolation=3),
T.RandomHorizontalFlip(p=0.5),
T.Pad(10),
T.RandomCrop((height, width)),
T.ToTensor(),
normalizer,
# T.RandomErasing(EPSILON=re),
T.RandomErasing(probability=0.4, mean=[0.485, 0.456, 0.406])
])
test_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.ToTensor(),
normalizer,
])
'''
num_instances=4
rmgs_flag = num_instances > 0
if rmgs_flag:
sampler = RandomIdentitySampler(dataset.target_train, num_instances)
else:
sampler = None
'''
source_train_loader = DataLoader(
Preprocessor(dataset.source_train, root=osp.join(dataset.source_images_dir, dataset.source_train_path),
transform=train_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=True, pin_memory=True, drop_last=True)
num_instances=0
rmgs_flag = num_instances > 0
if rmgs_flag:
sampler = RandomIdentitySampler(dataset.target_train, num_instances)
else:
sampler = None
target_train_loader = DataLoader(
UnsupervisedCamStylePreprocessor(dataset.target_train,
root=osp.join(dataset.target_images_dir, dataset.target_train_path),
camstyle_root=osp.join(dataset.target_images_dir,
dataset.target_train_camstyle_path),
num_cam=dataset.target_num_cam, transform=train_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=True, pin_memory=True, drop_last=True)
query_loader = DataLoader(
Preprocessor(dataset.query,
root=osp.join(dataset.target_images_dir, dataset.query_path), transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
query_loader_2 = DataLoader(
Preprocessor(dataset_2.query,
root=osp.join(dataset_2.target_images_dir, dataset_2.query_path), transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
gallery_loader = DataLoader(
Preprocessor(dataset.gallery,
root=osp.join(dataset.target_images_dir, dataset.gallery_path), transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
gallery_loader_2 = DataLoader(
Preprocessor(dataset_2.gallery,
root=osp.join(dataset_2.target_images_dir, dataset_2.gallery_path), transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
return dataset,dataset_2, num_classes, source_train_loader, target_train_loader, query_loader, gallery_loader, query_loader_2, gallery_loader_2
def main(args):
cudnn.deterministic = False
cudnn.benchmark = True
exp_database_dir = osp.join(args.exp_dir, string.capwords(args.dataset))
output_dir = osp.join(exp_database_dir, args.method, args.sub_method)
log_file = osp.join(output_dir, 'log.txt')
# Redirect print to both console and log file
sys.stdout = Logger(log_file)
# Create data loaders
dataset, num_classes, train_loader, _, _ = \
get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.combine_all,
args.workers, args.test_fea_batch)
# Create model
#model = seTest.resnst50().cuda()
model = resmap.create(args.arch,
ibn_type=args.ibn,
final_layer=args.final_layer,
neck=args.neck).cuda()
num_features = model.num_features
#num_features = 64
# print(model)
# print('\n')
feamap_factor = {'layer2': 8, 'layer3': 16, 'layer4': 32}
hei = args.height // feamap_factor[args.final_layer]
wid = args.width // feamap_factor[args.final_layer]
matcher = QAConv(num_features, hei, wid).cuda()
for arg in sys.argv:
print('%s ' % arg, end='')
print('\n')
# Criterion
criterion = ClassMemoryLoss(matcher, num_classes, num_features, hei, wid,
args.mem_batch_size).cuda()
# Optimizer
base_param_ids = set(map(id, model.base.parameters()))
new_params = [p for p in model.parameters() if id(p) not in base_param_ids]
param_groups = [{
'params': model.base.parameters(),
'lr': 0.1 * args.lr
}, {
'params': new_params,
'lr': args.lr
}, {
'params': criterion.parameters(),
'lr': args.lr
}]
optimizer = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=0.9,
weight_decay=5e-4,
nesterov=True)
# Load from checkpoint
start_epoch = 0
if args.resume or args.evaluate:
print('Loading checkpoint...')
if args.resume and (args.resume != 'ori'):
checkpoint = load_checkpoint(args.resume)
else:
checkpoint = load_checkpoint(
osp.join(output_dir, 'checkpoint.pth.tar'))
model.load_state_dict(checkpoint['model'])
criterion.load_state_dict(checkpoint['criterion'])
optimizer.load_state_dict(checkpoint['optim'])
start_epoch = checkpoint['epoch']
print("=> Start epoch {} ".format(start_epoch))
elif args.pre_epochs > 0:
pre_tr = PreTrainer(model, criterion, optimizer, train_loader,
args.pre_epochs, args.max_steps, args.num_trials)
result_file = osp.join(exp_database_dir, args.method,
'pretrain_metric.txt')
model, criterion, optimizer = pre_tr.train(result_file, args.method,
args.sub_method)
# Decay LR by a factor of 0.1 every step_size epochs
lr_scheduler = StepLR(optimizer,
step_size=args.step_size,
gamma=0.1,
last_epoch=start_epoch - 1)
model = nn.DataParallel(model).cuda()
criterion = nn.DataParallel(criterion).cuda()
enhance_data_aug = False
if not args.evaluate:
# Trainer
trainer = Trainer(model, criterion)
t0 = time.time()
# Start training
for epoch in range(start_epoch, args.epochs):
loss, acc = trainer.train(epoch, train_loader, optimizer)
lr = list(map(lambda group: group['lr'], optimizer.param_groups))
lr_scheduler.step()
train_time = time.time() - t0
print(
'* Finished epoch %d at lr=[%g, %g, %g]. Loss: %.3f. Acc: %.2f%%. Training time: %.0f seconds. \n'
%
(epoch + 1, lr[0], lr[1], lr[2], loss, acc * 100, train_time))
save_checkpoint(
{
'model': model.module.state_dict(),
'criterion': criterion.module.state_dict(),
'optim': optimizer.state_dict(),
'epoch': epoch + 1,
},
fpath=osp.join(output_dir, 'checkpoint.pth.tar'))
if not enhance_data_aug and epoch < args.epochs - 1 and acc > args.acc_thr:
enhance_data_aug = True
print('\nAcc = %.2f%% > %.2f%%. Start to Flip and Block.\n' %
(acc * 100, args.acc_thr * 100))
train_transformer = T.Compose([
T.Resize((args.height, args.width), interpolation=3),
T.Pad(10),
T.RandomCrop((args.height, args.width)),
T.RandomHorizontalFlip(0.5),
T.RandomRotation(5),
T.ColorJitter(brightness=(0.5, 2.0),
contrast=(0.5, 2.0),
saturation=(0.5, 2.0),
hue=(-0.1, 0.1)),
T.RandomOcclusion(args.min_size, args.max_size),
T.ToTensor(),
])
train_loader = DataLoader(Preprocessor(
dataset.train,
root=osp.join(dataset.images_dir, dataset.train_path),
transform=train_transformer),
batch_size=args.batch_size,
num_workers=args.workers,
shuffle=True,
pin_memory=True,
drop_last=True)
# Final test
print('Evaluate the learned model:')
t0 = time.time()
# Evaluator
evaluator = Evaluator(model)
avg_rank1 = 0
avg_mAP = 0
num_testsets = 0
results = {}
test_names = args.testset.strip().split(',')
for test_name in test_names:
if test_name not in datasets.names():
print('Unknown dataset: %s.' % test_name)
continue
testset, test_query_loader, test_gallery_loader = \
get_test_data(test_name, args.data_dir, args.height, args.width, args.workers, args.test_fea_batch)
if not args.do_tlift:
testset.has_time_info = False
test_rank1, test_mAP, test_rank1_rerank, test_mAP_rerank, test_rank1_tlift, test_mAP_tlift, test_dist, \
test_dist_rerank, test_dist_tlift, pre_tlift_dict = \
evaluator.evaluate(matcher, testset, test_query_loader, test_gallery_loader,
args.test_gal_batch, args.test_prob_batch,
args.tau, args.sigma, args.K, args.alpha)
results[test_name] = [test_rank1, test_mAP]
if test_name != args.dataset:
avg_rank1 += test_rank1
avg_mAP += test_mAP
num_testsets += 1
if testset.has_time_info:
print(
' %s: rank1=%.1f, mAP=%.1f, rank1_rerank=%.1f, mAP_rerank=%.1f,'
' rank1_rerank_tlift=%.1f, mAP_rerank_tlift=%.1f.\n' %
(test_name, test_rank1 * 100, test_mAP * 100,
test_rank1_rerank * 100, test_mAP_rerank * 100,
test_rank1_tlift * 100, test_mAP_tlift * 100))
else:
print(' %s: rank1=%.1f, mAP=%.1f.\n' %
(test_name, test_rank1 * 100, test_mAP * 100))
result_file = osp.join(exp_database_dir, args.method,
test_name + '_results.txt')
with open(result_file, 'a') as f:
f.write('%s/%s:\n' % (args.method, args.sub_method))
if testset.has_time_info:
f.write(
'\t%s: rank1=%.1f, mAP=%.1f, rank1_rerank=%.1f, mAP_rerank=%.1f, rank1_rerank_tlift=%.1f, '
'mAP_rerank_tlift=%.1f.\n\n' %
(test_name, test_rank1 * 100, test_mAP * 100,
test_rank1_rerank * 100, test_mAP_rerank * 100,
test_rank1_tlift * 100, test_mAP_tlift * 100))
else:
f.write('\t%s: rank1=%.1f, mAP=%.1f.\n\n' %
(test_name, test_rank1 * 100, test_mAP * 100))
if args.save_score:
test_gal_list = np.array(
[fname for fname, _, _, _ in testset.gallery], dtype=np.object)
test_prob_list = np.array(
[fname for fname, _, _, _ in testset.query], dtype=np.object)
test_gal_ids = [pid for _, pid, _, _ in testset.gallery]
test_prob_ids = [pid for _, pid, _, _ in testset.query]
test_gal_cams = [c for _, _, c, _ in testset.gallery]
test_prob_cams = [c for _, _, c, _ in testset.query]
test_score_file = osp.join(exp_database_dir, args.method,
args.sub_method,
'%s_score.mat' % test_name)
sio.savemat(test_score_file, {
'score': 1. - test_dist,
'score_rerank': 1. - test_dist_rerank,
'score_tlift': 1. - test_dist_tlift,
'gal_time': pre_tlift_dict['gal_time'],
'prob_time': pre_tlift_dict['prob_time'],
'gal_list': test_gal_list,
'prob_list': test_prob_list,
'gal_ids': test_gal_ids,
'prob_ids': test_prob_ids,
'gal_cams': test_gal_cams,
'prob_cams': test_prob_cams
},
oned_as='column',
do_compression=True)
test_time = time.time() - t0
avg_rank1 /= num_testsets
avg_mAP /= num_testsets
for key in results.keys():
print('%s: rank1=%.1f%%, mAP=%.1f%%.' %
(key, results[key][0] * 100, results[key][1] * 100))
print('Average: rank1=%.2f%%, mAP=%.2f%%.\n\n' %
(avg_rank1 * 100, avg_mAP * 100))
result_file = osp.join(exp_database_dir, args.method,
args.sub_method[:-5] + '_avg_results.txt')
with open(result_file, 'a') as f:
f.write('%s/%s:\n' % (args.method, args.sub_method))
if not args.evaluate:
f.write('\t Loss: %.3f, acc: %.2f%%. ' % (loss, acc * 100))
f.write("Train: %.0fs. " % train_time)
f.write("Test: %.0fs. " % test_time)
f.write('Rank1: %.2f%%, mAP: %.2f%%.\n' %
(avg_rank1 * 100, avg_mAP * 100))
for key in results.keys():
f.write('\t %s: Rank1: %.1f%%, mAP: %.1f%%.\n' %
(key, results[key][0] * 100, results[key][1] * 100))
f.write('\n')
if not args.evaluate:
print('Finished training at epoch %d, loss = %.3f, acc = %.2f%%.\n' %
(epoch + 1, loss, acc * 100))
print(
"Total training time: %.3f sec. Average training time per epoch: %.3f sec."
% (train_time, train_time / (args.epochs - start_epoch + 1)))
print("Total testing time: %.3f sec.\n" % test_time)
for arg in sys.argv:
print('%s ' % arg, end='')
print('\n')
def main(args):
np.random.seed(args.seed) # With the seed reset (every time), the same set of numbers will appear every time.
torch.manual_seed(args.seed) # Sets the seed for generating random numbers.
cudnn.benchmark = True # This flag allows you to enable the inbuilt cudnn auto-tuner to find the best algorithm to use for your hardware. It enables benchmark mode in cudnn.
# Redirect print to both console and log file
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
print(args)
shutil.copy(sys.argv[0], osp.join(args.logs_dir,
osp.basename(sys.argv[0])))
# Create data loaders
if args.height is None or args.width is None:
args.height, args.width = (256, 128)
dataset, num_classes, train_loader, val_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.height,
args.width, args.batch_size * 8, args.workers, #https://deeplizard.com/learn/video/kWVgvsejXsE#:~:text=The%20num_workers%20attribute%20tells%20the,sequentially%20inside%20the%20main%20process.
)
# Create model
model = models.create("ft_net_inter",
num_classes=num_classes, stride=args.stride)
# Load from checkpoint
start_epoch = 0
best_top1 = 0
top1 = 0
is_best = False
if args.checkpoint is not None:
if args.evaluate:
checkpoint = load_checkpoint(args.checkpoint)
param_dict = model.state_dict() # A state_dict is simply a Python dictionary object that maps each layer to its parameter tensor.
for k, v in checkpoint['state_dict'].items():
if 'model' in k:
param_dict[k] = v
model.load_state_dict(param_dict)
else:
model.model.load_param(args.checkpoint)
model = model.cuda()
# Distance metric
metric = None
# Evaluator
evaluator = Evaluator(model, use_cpu=args.use_cpu)
if args.evaluate:
print("Test:")
evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
return
train_transformer = [
T.Resize((args.height, args.width), interpolation=3),
T.RandomHorizontalFlip(),
T.Pad(10),
T.RandomCrop((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
T.RandomErasing(probability=0.5),
]
train_transformer = T.Compose(train_transformer)
for cluster_epoch in range(args.cluster_epochs):
# -------------------------Stage 1 intra camera training--------------------------
# Cluster and generate new dataset and model
# Divides the training set into (subsets) and according to that each camera id is there for each image
# then it forms clustering on each subset according to the pair wise similarity
# then assigning images with in each cluster with identical label
# then cross entropy loss is used
cluster_result = get_intra_cam_cluster_result(model, train_loader,
args.class_number_stage1,
args.linkage)
cluster_datasets = [
datasets.create("cluster", osp.join(args.data_dir, args.dataset),
cluster_result[cam_id], cam_id)
for cam_id in cluster_result.keys()
]
cluster_dataloaders = [
DataLoader(Preprocessor(dataset.train_set,
root=dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size,
num_workers=args.workers,
shuffle=True,
pin_memory=False,
drop_last=True) for dataset in cluster_datasets
]
param_dict = model.model.state_dict()
model = models.create("ft_net_intra",
num_classes=[
args.class_number_stage1
for cam_id in cluster_result.keys()
],
stride=args.stride)
model_param_dict = model.model.state_dict()
for k, v in model_param_dict.items():
if k in param_dict.keys():
model_param_dict[k] = param_dict[k]
model.model.load_state_dict(model_param_dict)
model = model.cuda()
criterion = nn.CrossEntropyLoss().cuda()
# Optimizer
param_groups = make_params(model, args.lr, args.weight_decay)
optimizer = torch.optim.SGD(param_groups, momentum=0.9)
# Trainer
trainer = IntraCameraTrainer(
model, criterion, warm_up_epoch=args.warm_up)
print("start training")
# Start training
for epoch in range(0, args.epochs_stage1):
trainer.train(cluster_epoch,
epoch,
cluster_dataloaders,
optimizer,
print_freq=args.print_freq,
)
# -------------------------------------------Stage 2 inter camera training-----------------------------------
mix_rate = get_mix_rate(
args.mix_rate, cluster_epoch, args.cluster_epochs, power=args.decay_factor)
cluster_result = get_inter_cam_cluster_result(
model,
train_loader,
args.class_number_stage2,
args.linkage,
mix_rate,
use_cpu=args.use_cpu)
cluster_dataset = datasets.create(
"cluster", osp.join(args.data_dir, args.dataset), cluster_result,
0)
cluster_dataloaders = DataLoader(
Preprocessor(cluster_dataset.train_set,
root=cluster_dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size_stage2,
num_workers=args.workers,
sampler=RandomIdentitySampler(cluster_dataset.train_set,
args.batch_size_stage2,
args.instances),
pin_memory=False,
drop_last=True)
param_dict = model.model.state_dict()
model = models.create("ft_net_inter",
num_classes=args.class_number_stage2,
stride=args.stride)
model.model.load_state_dict(param_dict)
model = model.cuda()
# Criterion
criterion_entropy = nn.CrossEntropyLoss().cuda()
criterion_triple = TripletLoss(margin=args.margin).cuda()
# Optimizer
param_groups = make_params(model,
args.lr * args.batch_size_stage2 / 32,
args.weight_decay)
optimizer = torch.optim.SGD(param_groups, momentum=0.9)
# Trainer
trainer = InterCameraTrainer(model,
criterion_entropy,
criterion_triple,
warm_up_epoch=args.warm_up,
)
print("start training")
# Start training
for epoch in range(0, args.epochs_stage2):
trainer.train(cluster_epoch,
epoch,
cluster_dataloaders,
optimizer,
print_freq=args.print_freq)
if (cluster_epoch + 1) % 5 == 0: # in 4th, 9th, 14th epochs, see in the output
evaluator = Evaluator(model, use_cpu=args.use_cpu)
top1, mAP = evaluator.evaluate(
test_loader, dataset.query, dataset.gallery, metric, return_mAP=True)
is_best = top1 > best_top1
best_top1 = max(top1, best_top1)
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': cluster_epoch + 1,
'best_top1': best_top1,
'cluster_epoch': cluster_epoch + 1,
},
is_best,
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
if cluster_epoch == (args.cluster_epochs - 1):
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': cluster_epoch + 1,
'best_top1': best_top1,
'cluster_epoch': cluster_epoch + 1,
},
False,
fpath=osp.join(args.logs_dir, 'latest.pth.tar'))
print('\n * cluster_epoch: {:3d} top1: {:5.1%} best: {:5.1%}{}\n'.
format(cluster_epoch, top1, best_top1, ' *' if is_best else ''))
# Final test
print('Test with best model:')
checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
model.load_state_dict(checkpoint['state_dict'])
best_rank1, mAP = evaluator.evaluate(
test_loader, dataset.query, dataset.gallery, metric, return_mAP=True)
