def get_optim_params(config, params, device):
if config.loss_name not in ['softmax', 'triplet', 'weight_softmax']:
raise ValueError('wrong loss name')
if config.loss_name is 'triplet':
criterion = TripletLoss(margin=config.margin)
optimizer = torch.optim.Adam(params,
lr=config.lr,
weight_decay=config.weight_decay)
return criterion, optimizer
optimizer = torch.optim.SGD(params,
momentum=config.momentum,
weight_decay=config.weight_decay,
nesterov=True)
if config.loss_name is 'softmax':
criterion = nn.CrossEntropyLoss()
else:
criterion = SoftCrossEntropyLoss()
criterion = criterion.to(device)
return criterion, optimizer
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)
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
if args.loss == 'triplet':
assert args.num_instances > 1, 'TripletLoss requires num_instances > 1'
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
dataset, num_classes, train_loader, val_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir,
args.batch_size, args.workers, args.num_instances,
combine_trainval=args.combine_trainval)
# Create model
if args.loss == 'xentropy':
model = InceptionNet(num_classes=num_classes,
num_features=args.features,
dropout=args.dropout)
elif args.loss == 'oim':
model = InceptionNet(num_features=args.features,
norm=True,
dropout=args.dropout)
elif args.loss == 'triplet':
model = InceptionNet(num_features=args.features, dropout=args.dropout)
else:
raise ValueError("Cannot recognize loss type:", args.loss)
model = torch.nn.DataParallel(model).cuda()
# Load from checkpoint
if args.resume:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch']
best_top1 = checkpoint['best_top1']
print("=> start epoch {} best top1 {:.1%}".format(
args.start_epoch, best_top1))
else:
best_top1 = 0
# Distance metric
metric = DistanceMetric(algorithm=args.dist_metric)
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
metric.train(model, train_loader)
print("Validation:")
evaluator.evaluate(val_loader, dataset.val, dataset.val, metric)
print("Test:")
evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
return
# Criterion
if args.loss == 'xentropy':
criterion = torch.nn.CrossEntropyLoss()
elif args.loss == 'oim':
criterion = OIMLoss(model.module.num_features,
num_classes,
scalar=args.oim_scalar,
momentum=args.oim_momentum)
elif args.loss == 'triplet':
criterion = TripletLoss(margin=args.triplet_margin)
else:
raise ValueError("Cannot recognize loss type:", args.loss)
criterion.cuda()
# Optimizer
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
else:
raise ValueError("Cannot recognize optimizer type:", args.optimizer)
# Trainer
trainer = Trainer(model, criterion)
# Schedule learning rate
def adjust_lr(epoch):
if args.optimizer == 'sgd':
lr = args.lr * (0.1**(epoch // 60))
elif args.optimizer == 'adam':
lr = args.lr if epoch <= 100 else \
args.lr * (0.001 ** (epoch - 100) / 50)
else:
raise ValueError("Cannot recognize optimizer type:",
args.optimizer)
for g in optimizer.param_groups:
g['lr'] = lr
# Start training
for epoch in range(args.start_epoch, args.epochs):
adjust_lr(epoch)
trainer.train(epoch, train_loader, optimizer)
top1 = evaluator.evaluate(val_loader, dataset.val, dataset.val)
is_best = top1 > best_top1
best_top1 = max(top1, best_top1)
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': epoch + 1,
'best_top1': best_top1,
},
is_best,
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'.
format(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'])
metric.train(model, train_loader)
evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
def main(args):
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
print(args)
# Create data loaders
dataset, num_classes, source_train_loader, query_loader_2, gallery_loader_2, query_loader_3, gallery_loader_3 = \
get_data(args.data_dir, args.source, args.target, args.height,
args.width, args.batch_size, args.num_instance, args.workers)
# Create model
MaskNet, TaskNet = models.create(args.arch,
num_features=args.features,
dropout=args.dropout,
num_classes=num_classes)
# Load from checkpoint
start_epoch = 0
if args.resume:
checkpoint = load_checkpoint(args.resume)
MaskNet.load_state_dict(checkpoint['MaskNet'])
TaskNet.load_state_dict(checkpoint['TaskNet'])
start_epoch = checkpoint['epoch']
print("=> Start epoch {} ".format(start_epoch))
MaskNet = nn.DataParallel(MaskNet).cuda()
TaskNet = nn.DataParallel(TaskNet).cuda()
# Evaluator
evaluator = Evaluator([MaskNet, TaskNet])
if args.evaluate:
print("Test:")
print("partial-iLIDS:")
evaluator.evaluate(query_loader_2, gallery_loader_2, dataset.query_2,
dataset.gallery_2, args.output_feature)
print("partial-REID:")
evaluator.evaluate(query_loader_3, gallery_loader_3, dataset.query_3,
dataset.gallery_3, args.output_feature)
return
# Criterion
criterion = []
criterion.append(nn.CrossEntropyLoss().cuda())
criterion.append(TripletLoss(margin=args.margin))
criterion.append(nn.MSELoss(reduce=True, size_average=True).cuda())
# Optimizer
param_groups = [
{
'params': MaskNet.module.parameters(),
'lr_mult': 0.1
},
]
optimizer_Mask = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
#
base_param_ids = set(map(id, TaskNet.module.base.parameters()))
new_params = [
p for p in TaskNet.parameters() if id(p) not in base_param_ids
]
param_groups = [{
'params': TaskNet.module.base.parameters(),
'lr_mult': 0.1
}, {
'params': new_params,
'lr_mult': 1.0
}]
optimizer_Ide = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
# Trainer
trainer = Trainer([MaskNet, TaskNet], criterion)
# Schedule learning rate
def adjust_lr(epoch):
step_size = 10
if epoch <= 9:
lr = 0.0008 * (epoch / 10.0)
elif epoch <= 16:
lr = 0.1
elif epoch <= 23:
lr = 0.001
else:
lr = 0.0001
for g in optimizer_Mask.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
for g in optimizer_Ide.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
tmp_2 = best_2 = 0
tmp_3 = best_3 = 0
# Start training
for epoch in range(start_epoch, args.epochs):
adjust_lr(epoch)
trainer.train(epoch, [source_train_loader],
[optimizer_Mask, optimizer_Ide], args.batch_size)
save_checkpoint(
{
'MaskNet': MaskNet.module.state_dict(),
'TaskNet': TaskNet.module.state_dict(),
'epoch': epoch + 1,
},
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
if epoch == 9:
save_checkpoint(
{
'MaskNet': MaskNet.module.state_dict(),
'TaskNet': TaskNet.module.state_dict(),
'epoch': epoch + 1,
},
fpath=osp.join(args.logs_dir, 'epoch9_checkpoint.pth.tar'))
evaluator = Evaluator([MaskNet, TaskNet])
if epoch > 9 and epoch % 1 == 0:
tmp_2 = evaluator.evaluate(query_loader_2, gallery_loader_2,
dataset.query_2, dataset.gallery_2,
args.output_feature)
tmp_3 = evaluator.evaluate(query_loader_3, gallery_loader_3,
dataset.query_3, dataset.gallery_3,
args.output_feature)
if (tmp_2 > best_2):
save_checkpoint(
{
'MaskNet': MaskNet.module.state_dict(),
'TaskNet': TaskNet.module.state_dict(),
'epoch': epoch + 1,
},
fpath=osp.join(args.logs_dir,
'best_checkpoint_piLIDS.pth.tar'))
best_2 = tmp_2
print("iLIDS_best:", best_2)
#
if (tmp_3 > best_3):
save_checkpoint(
{
'MaskNet': MaskNet.module.state_dict(),
'TaskNet': TaskNet.module.state_dict(),
'epoch': epoch + 1,
},
fpath=osp.join(args.logs_dir, 'best_checkpoint_pREID.pth.tar'))
best_3 = tmp_3
print("REID_best:", best_3)
print('\n * Finished epoch {:3d} \n'.format(epoch))
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch == 'inception' else \
(256, 128)
train, val, trainval = [], [], []
numbers = [0, 0, 0]
dataset_cuhk03 = merge('cuhk03', train, val, trainval, numbers, args.data_dir, args.split)
dataset_market1501 = merge('market1501', train, val, trainval, numbers, args.data_dir, args.split)
merge('cuhksysu', train, val, trainval, numbers, args.data_dir, args.split)
merge('mars', train, val, trainval, numbers, args.data_dir, args.split)
num_train_ids, num_val_ids, num_trainval_ids = numbers
assert num_val_ids == dataset_cuhk03.num_val_ids + dataset_market1501.num_val_ids
print("============================================")
print("JSTL dataset loaded")
print(" subset | # ids | # images")
print(" ---------------------------")
print(" train | {:5d} | {:8d}"
.format(num_train_ids, len(train)))
print(" val | {:5d} | {:8d}"
.format(num_val_ids, len(val)))
print(" trainval | {:5d} | {:8d}"
.format(num_trainval_ids, len(trainval)))
query_cuhk03, gallery_cuhk03 = dataset_cuhk03.query, dataset_cuhk03.gallery
query_market1501, gallery_market1501 = dataset_market1501.query, dataset_market1501.gallery
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_set = trainval if args.combine_trainval else train
num_classes = (num_trainval_ids if args.combine_trainval
else num_train_ids)
train_transformer = T.Compose([
T.RandomSizedRectCrop(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
normalizer,
])
test_transformer = T.Compose([
T.RectScale(args.height, args.width),
T.ToTensor(),
normalizer,
])
train_loader = DataLoader(
Preprocessor(train_set, root=args.data_dir,
transform=train_transformer),
batch_size=args.batch_size, num_workers=args.workers,
sampler=RandomIdentitySampler(train_set, args.num_instances),
pin_memory=True, drop_last=True)
val_loader = DataLoader(
Preprocessor(val, root=args.data_dir,
transform=test_transformer),
batch_size=args.batch_size, num_workers=args.workers,
shuffle=False, pin_memory=True)
test_loader_cuhk03 = DataLoader(
Preprocessor(list(set(query_cuhk03) | set(gallery_cuhk03)),
root=dataset_cuhk03.images_dir, transform=test_transformer),
batch_size=args.batch_size, num_workers=args.workers,
shuffle=False, pin_memory=True)
test_loader_market1501 = DataLoader(
Preprocessor(list(set(query_market1501) | set(gallery_market1501)),
root=dataset_market1501.images_dir, transform=test_transformer),
batch_size=args.batch_size, num_workers=args.workers,
shuffle=False, pin_memory=True)
# Create model
# Hacking here to let the classifier be the last feature embedding layer
# Net structure: avgpool -> FC(1024) -> FC(args.features)
model = models.create(args.arch, num_features=1024,
dropout=args.dropout, num_classes=args.features)
# Load from checkpoint
start_epoch = best_top1 = 0
if args.resume:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
best_top1 = checkpoint['best_top1']
print("=> Start epoch {} best top1 {:.1%}"
.format(start_epoch, best_top1))
model = nn.DataParallel(model).cuda()
# Distance metric
metric = DistanceMetric(algorithm=args.dist_metric)
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
metric.train(model, train_loader)
print("Validation:")
evaluator.evaluate(val_loader, val, val, metric)
print("Test(cuhk03):")
evaluator.evaluate(test_loader_cuhk03, query_cuhk03, gallery_cuhk03, metric)
print("Test(market1501):")
evaluator.evaluate(test_loader_market1501, query_market1501, gallery_market1501, metric)
return
# Criterion
criterion = TripletLoss(margin=args.margin).cuda()
# Optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr,
weight_decay=args.weight_decay)
# Trainer
trainer = Trainer(model, criterion)
# Schedule learning rate
def adjust_lr(epoch):
lr = args.lr if epoch <= 100 else \
args.lr * (0.001 ** ((epoch - 100) / 50.0))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
# Start training
for epoch in range(start_epoch, args.epochs):
adjust_lr(epoch)
trainer.train(epoch, train_loader, optimizer)
if epoch < args.start_save:
continue
top1 = evaluator.evaluate(val_loader, val, val)
is_best = top1 > best_top1
best_top1 = max(top1, best_top1)
save_checkpoint({
'state_dict': model.module.state_dict(),
'epoch': epoch + 1,
'best_top1': best_top1,
}, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'.
format(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.module.load_state_dict(checkpoint['state_dict'])
metric.train(model, train_loader)
print("Test(cuhk03):")
evaluator.evaluate(test_loader_cuhk03, query_cuhk03, gallery_cuhk03, metric)
print("Test(market1501):")
evaluator.evaluate(test_loader_market1501, query_market1501, gallery_market1501, metric)
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch == 'inception' else \
(256, 128)
# get source data
src_dataset, src_extfeat_loader = \
get_source_data(args.src_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# get target data
tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \
get_data(args.tgt_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
# Hacking here to let the classifier be the number of source ids
if args.src_dataset == 'dukemtmc':
model = models.create(args.arch, num_classes=632, pretrained=False)
coModel = models.create(args.arch, num_classes=632, pretrained=False)
elif args.src_dataset == 'market1501':
model = models.create(args.arch, num_classes=676, pretrained=False)
coModel = models.create(args.arch, num_classes=676, pretrained=False)
else:
raise RuntimeError(
'Please specify the number of classes (ids) of the network.')
# Load from checkpoint
if args.resume:
print(
'Resuming checkpoints from finetuned model on another dataset...\n'
)
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'], strict=False)
coModel.load_state_dict(checkpoint['state_dict'], strict=False)
else:
raise RuntimeWarning('Not using a pre-trained model.')
model = nn.DataParallel(model).cuda()
coModel = nn.DataParallel(coModel).cuda()
# evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery)
# if args.evaluate: return
# Criterion
criterion = [
TripletLoss(args.margin,
args.num_instances,
isAvg=False,
use_semi=False).cuda(),
TripletLoss(args.margin,
args.num_instances,
isAvg=False,
use_semi=False).cuda(),
]
# Optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
coOptimizer = torch.optim.Adam(coModel.parameters(), lr=args.lr)
optims = [optimizer, coOptimizer]
# training stage transformer on input images
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transformer = T.Compose([
T.Resize((args.height, args.width)),
T.RandomHorizontalFlip(),
T.ToTensor(), normalizer,
T.RandomErasing(probability=0.5, sh=0.2, r1=0.3)
])
# # Start training
for iter_n in range(args.iteration):
if args.lambda_value == 0:
source_features = 0
else:
# get source datas' feature
source_features, _ = extract_features(model,
src_extfeat_loader,
print_freq=args.print_freq)
# synchronization feature order with src_dataset.train
source_features = torch.cat([
source_features[f].unsqueeze(0)
for f, _, _ in src_dataset.train
], 0)
# extract training images' features
print('Iteration {}: Extracting Target Dataset Features...'.format(
iter_n + 1))
target_features, tarNames = extract_features(
model, tgt_extfeat_loader, print_freq=args.print_freq)
# synchronization feature order with dataset.train
target_features = torch.cat([
target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval
], 0)
target_real_label = np.asarray(
[tarNames[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval])
numTarID = len(set(target_real_label))
# calculate distance and rerank result
print('Calculating feature distances...')
target_features = target_features.numpy()
cluster = KMeans(n_clusters=numTarID, n_jobs=8, n_init=1)
# select & cluster images as training set of this epochs
print('Clustering and labeling...')
clusterRes = cluster.fit(target_features)
labels, centers = clusterRes.labels_, clusterRes.cluster_centers_
labels = splitLowconfi(target_features, labels, centers)
# num_ids = len(set(labels))
# print('Iteration {} have {} training ids'.format(iter_n+1, num_ids))
# generate new dataset
new_dataset, unknown_dataset = [], []
# assign label for target ones
unknownLab = labelNoise(torch.from_numpy(target_features),
torch.from_numpy(labels))
# unknownFeats = target_features[labels==-1,:]
unCounter = 0
for (fname, _, cam), label in zip(tgt_dataset.trainval, labels):
if label == -1:
unknown_dataset.append(
(fname, int(unknownLab[unCounter]), cam)) # unknown data
unCounter += 1
continue
# dont need to change codes in trainer.py _parsing_input function and sampler function after add 0
new_dataset.append((fname, label, cam))
print('Iteration {} have {} training images'.format(
iter_n + 1, len(new_dataset)))
train_loader = DataLoader(Preprocessor(new_dataset,
root=tgt_dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size,
num_workers=4,
sampler=RandomIdentitySampler(
new_dataset, args.num_instances),
pin_memory=True,
drop_last=True)
# hard samples
unLoader = DataLoader(Preprocessor(unknown_dataset,
root=tgt_dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size,
num_workers=4,
sampler=RandomIdentitySampler(
unknown_dataset, args.num_instances),
pin_memory=True,
drop_last=True)
# train model with new generated dataset
trainer = CoTrainerAsy(model, coModel, train_loader, unLoader,
criterion, optims)
# trainer = CoTeaching(
# model, coModel, train_loader, unLoader, criterion, optims
# )
# trainer = CoTrainerAsySep(
# model, coModel, train_loader, unLoader, criterion, optims
# )
evaluator = Evaluator(model, print_freq=args.print_freq)
#evaluatorB = Evaluator(coModel, print_freq=args.print_freq)
# Start training
for epoch in range(args.epochs):
trainer.train(epoch,
remRate=0.2 + (0.6 / args.iteration) *
(1 + iter_n)) # to at most 80%
# trainer.train(epoch, remRate=0.7+(0.3/args.iteration)*(1+iter_n))
# test only
rank_score = evaluator.evaluate(test_loader, tgt_dataset.query,
tgt_dataset.gallery)
#print('co-model:\n')
#rank_score = evaluatorB.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery)
# Evaluate
rank_score = evaluator.evaluate(test_loader, tgt_dataset.query,
tgt_dataset.gallery)
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch + 1,
'best_top1': rank_score.market1501[0],
},
True,
fpath=osp.join(args.logs_dir, 'adapted.pth.tar'))
return (rank_score.map, rank_score.market1501[0])
def main(args):
print(args)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
if args.big_height is None or args.big_width is None or args.target_height is None or args.target_width is None:
args.big_height, args.big_width, args.target_height, args.target_width = (
256, 256, 224, 224)
dataset, num_classes, train_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.big_height, args.big_width,
args.target_height, args.target_width, args.batch_size, args.num_instances,
args.workers, args.combine_trainval)
# Create model
model = models.create(args.arch,
num_classes=num_classes,
num_features=args.features)
# print(model)
# Load from checkpoint
start_epoch = best = 0
if args.weights and args.arch != 'dpn107' and args.arch != 'densenet161':
checkpoint = load_checkpoint(args.weights)
if args.arch == 'cross_trihard_senet101' or args.arch == 'cross_trihard_se_resnet152':
del (checkpoint['last_linear.weight'])
del (checkpoint['last_linear.bias'])
model.base.load_state_dict(checkpoint)
#model.base.load_param(args.weights)
elif args.arch == 'Hrnet48':
pass
elif args.arch == 'se_152_ibn':
del (checkpoint['last_linear.weight'])
del (checkpoint['last_linear.bias'])
model.base.load_state_dict(checkpoint, strict=False)
elif args.arch == 'densenet169_ibn_a':
checkpoint = torch.load(args.weights)['state_dict']
del (checkpoint['module.classifier.weight'])
del (checkpoint['module.classifier.bias'])
model.load_state_dict(
{k.replace('module.', ''): v
for k, v in checkpoint.items()},
strict=False)
else:
del (checkpoint['fc.weight'])
del (checkpoint['fc.bias'])
model.base.load_state_dict(checkpoint)
if args.resume and not args.weights:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
print("=> Start epoch {} best top1 {:.1%}".format(start_epoch, best))
model = nn.DataParallel(model).cuda()
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
print("Test:")
evaluator.evaluate(test_loader, dataset.query, dataset.gallery)
return
# Criterion
ranking_loss = nn.MarginRankingLoss(margin=args.margin).cuda()
criterion = { 'crossentropy': nn.CrossEntropyLoss().cuda(), \
'trihard': TripletLoss(ranking_loss).cuda() }
# Optimizer
if hasattr(model.module, 'base'):
base_params = []
base_bn_params = []
for name, p in model.module.base.named_parameters():
if 'bn' in name:
base_bn_params.append(p)
else:
base_params.append(p)
base_param_ids = set(map(id, model.module.base.parameters()))
new_params = [
p for p in model.parameters() if id(p) not in base_param_ids
]
param_groups = [{
'params': base_params,
'lr_mult': 0.1
}, {
'params': base_bn_params,
'lr_mult': 0.1
}, {
'params': new_params,
'lr_mult': args.lr_mult
}]
else:
param_groups = model.parameters()
if args.optimizer == 0:
optimizer = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
else:
print('Adam')
optimizer = torch.optim.Adam(params=param_groups,
lr=args.lr,
weight_decay=args.weight_decay)
# Trainer
trainer = Cross_Trihard_Trainer(model,
criterion,
metric_loss_weight=args.metric_loss_weight)
# Schedule learning rate
def adjust_lr(epoch):
step_size, step_size2, step_size3 = args.step_size, args.step_size2, args.step_size3
#lr = args.lr * (0.1 ** (epoch // step_size))
if epoch <= step_size:
lr = args.lr
elif epoch <= step_size2:
lr = args.lr * 0.1
elif epoch <= step_size3:
lr = args.lr * 0.01
else:
lr = args.lr * 0.001
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
return lr
# Start training
for epoch in range(start_epoch + 1, args.epochs + 1):
lr = adjust_lr(epoch)
trainer.train(epoch,
train_loader,
optimizer,
lr,
warm_up=True,
warm_up_ep=args.warm_up_ep)
if epoch % args.epoch_inter == 0 or epoch >= args.dense_evaluate:
tmp_res = evaluator.evaluate(test_loader, dataset.query,
dataset.gallery)
print('tmp_res: ', tmp_res)
print('best: ', best)
if tmp_res > best and epoch >= args.start_save:
best = tmp_res
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch,
},
False,
fpath=osp.join(args.logs_dir, 'pass%d.pth.tar' % (epoch)))
def main(args):
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
dataset, num_classes, train_loader, query_loader, gallery_loader = \
get_data(args.dataset, args.data_dir, args.height,
args.width, args.batch_size, args.camstyle, args.re, args.num_instances, args.workers)
# Create model
model = models.create(args.arch,
num_features=2048,
dropout=args.dropout,
num_classes=args.features)
# Load from checkpoint
start_epoch = 0
if args.resume:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
print("=> Start epoch {} ".format(start_epoch))
model = nn.DataParallel(model).cuda()
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
print("Test:")
evaluator.evaluate(query_loader, gallery_loader, dataset.query,
dataset.gallery, args.output_feature, args.rerank)
return
# Criterion
# criterion = TripletLoss(margin=args.margin).cuda()
criterion_cro = nn.CrossEntropyLoss().cuda()
criterion_tri = TripletLoss(margin=args.margin).cuda()
# Optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# Trainer
trainer = Trainer(model, criterion_cro, criterion_tri)
# Schedule learning rate
def adjust_lr(epoch):
lr = args.lr if epoch <= 100 else \
args.lr * (0.001 ** ((epoch - 100) / 50.0))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
# Start training
for epoch in range(start_epoch, args.epochs):
adjust_lr(epoch)
trainer.train(epoch, train_loader, optimizer)
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch + 1,
},
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print('\n * Finished epoch {:3d} \n'.format(epoch))
# Final test
print('Test with best model:')
evaluator = Evaluator(model)
evaluator.evaluate(query_loader, gallery_loader, dataset.query,
dataset.gallery, args.output_feature, args.rerank)
def main(args):
print(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
if args.big_height is None or args.big_width is None or args.target_height is None or args.target_width is None:
args.big_height, args.big_width, args.target_height, args.target_width = (
256, 256, 224, 224)
dataset, num_classes, train_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.big_height, args.big_width,
args.target_height, args.target_width, args.batch_size, args.num_instances,
args.workers, args.combine_trainval)
# Create model
model = models.create(args.arch,
num_classes=num_classes,
num_features=args.features,
is_cls=args.is_cls)
# Load from checkpoint
start_epoch = best = 0
if args.weights:
#model_dict = model.state_dict()
#checkpoint_load = {k: v for k, v in (checkpoint['state_dict']).items() if k in model_dict}
#model_dict.update(checkpoint_load)
#model.load_state_dict(model_dict)
if args.arch == 'cross_trihard_senet101':
model.base.load_param(args.weights)
else:
checkpoint = load_checkpoint(args.weights)
del (checkpoint['fc.weight'])
del (checkpoint['fc.bias'])
model.base.load_state_dict(checkpoint)
if args.resume and not args.weights:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
print("=> Start epoch {} best top1 {:.1%}".format(start_epoch, best))
model = nn.DataParallel(model).cuda()
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
print("Test:")
evaluator.evaluate(test_loader, dataset.query, dataset.gallery)
return
# Criterion
ranking_loss = nn.MarginRankingLoss(margin=args.margin).cuda()
if args.multi_attribute == 1:
criterion = { 'MultiAttributeLoss': MultiAttributeLoss(is_cls=args.is_cls).cuda(), \
'trihard': TripletLoss(ranking_loss).cuda() }
else:
criterion = { 'TypeAttributeLoss': TypeAttributeLoss(is_cls=args.is_cls).cuda(), \
'trihard': TripletLoss(ranking_loss).cuda() }
# Optimizer
if hasattr(model.module, 'base'):
base_params = []
base_bn_params = []
for name, p in model.module.base.named_parameters():
if 'bn' in name:
base_bn_params.append(p)
else:
base_params.append(p)
base_param_ids = set(map(id, model.module.base.parameters()))
new_params = [
p for p in model.parameters() if id(p) not in base_param_ids
]
param_groups = [{
'params': base_params,
'lr_mult': 0.1
}, {
'params': base_bn_params,
'lr_mult': 0.1
}, {
'params': new_params,
'lr_mult': args.lr_mult
}]
else:
param_groups = model.parameters()
if args.optimizer == 0:
print('SGD')
optimizer = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
else:
print('Adam')
optimizer = torch.optim.Adam(params=param_groups,
lr=args.lr,
weight_decay=args.weight_decay)
# Trainer
if args.multi_attribute == 1:
print('Multi Attribute')
trainer = Multi_Attribute_Trainer(
model,
criterion,
metric_loss_weight=args.metric_loss_weight,
sub_task_loss_weight=args.sub_task_loss_weight)
else:
print('Type Attribute')
trainer = Type_Attribute_Trainer(
model,
criterion,
metric_loss_weight=args.metric_loss_weight,
sub_task_loss_weight=args.sub_task_loss_weight)
# Schedule learning rate
def adjust_lr(epoch):
step_size, step_size2, step_size3 = args.step_size, args.step_size2, args.step_size3
if epoch <= step_size:
lr = args.lr
elif epoch <= step_size2:
lr = args.lr * 0.1
elif epoch <= step_size3:
lr = args.lr * 0.01
else:
lr = args.lr * 0.001
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
return lr
# Start training
for epoch in range(start_epoch + 1, args.epochs + 1):
lr = adjust_lr(epoch)
trainer.train(epoch,
train_loader,
optimizer,
lr,
warm_up=True,
warm_up_ep=args.warm_up_ep)
if epoch % args.epoch_inter == 0 or epoch >= args.dense_evaluate:
tmp_res = evaluator.evaluate(test_loader,
dataset.query,
dataset.gallery,
is_attribute=True)
if tmp_res > best and epoch >= args.start_save:
best = tmp_res
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch,
},
False,
fpath=osp.join(args.logs_dir, 'pass%d.pth.tar' % (epoch)))
def main(args):
random.seed(args.seed)
np.random.seed(1)
torch.manual_seed(1)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# params
params = {}
if args.model_type == 'masks':
params['num_m_features'] = args.num_m_features
params['masks'] = args.masks
else:
print('unkrown model type.')
return
# Create data loaders
if args.height is None or args.width is None:
args.height, args.width = (256, 256)
dataset, num_classes, random_train_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.logs_dir, args.model_type, params,
args.height, args.width, args.crop_height, args.crop_width,
args.batch_size, args.workers, args.combine_trainval)
# Create model
model = models.create(args.arch, num_classes=num_classes, params=params)
# Load from checkpoint
start_epoch = best_top1 = best_mAP = 0
if args.weights:
checkpoint = load_checkpoint(args.weights)
model_dict = model.state_dict()
checkpoint_load = {
k: v
for k, v in (checkpoint['state_dict']).items() if k in model_dict
}
model_dict.update(checkpoint_load)
model.load_state_dict(model_dict)
if args.resume and not args.weights:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
print("=> Start epoch {} best top1 {:.1%}".format(
start_epoch, best_top1))
model = nn.DataParallel(model, [0, 1, 2, 3]).cuda()
# Criterion
criterion = TripletLoss().cuda()
# Optimizer
base_params = []
for name, p in model.module.base.named_parameters():
base_params.append(p)
base_param_ids = set(map(id, model.module.base.parameters()))
new_params = [p for p in model.parameters() if id(p) not in base_param_ids]
if args.model_type == 'masks':
param_groups = [{
'params': base_params,
'lr_mult': args.lrm
}, {
'params': new_params,
'lr_mult': 1.0
}]
else:
print('unkrown model type.')
return
optimizer = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
# Trainer
trainer = Trainer(model, criterion, num_classes, args.logs_dir)
# Evaluator
evaluator = Evaluator(model)
# Schedule learning rate
def adjust_lr(epoch):
step_size = args.step_size
if epoch < step_size:
lr = args.lr
elif epoch >= step_size and epoch < args.epochs:
lr = args.lr * 0.1
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
return lr
# Start training
for epoch in range(start_epoch, args.epochs):
lr = adjust_lr(epoch)
if epoch < args.warm_up_ep:
trainer.train(epoch, random_train_loader, optimizer, lr, True,
args.warm_up_ep)
else:
trainer.train(epoch, random_train_loader, optimizer, lr, False,
args.warm_up_ep)
if epoch < args.start_save:
continue
if epoch % 10 == 9:
print('Epoch: [%d]' % epoch)
top1, mAP = evaluator.evaluate(test_loader, dataset.query,
dataset.gallery)
is_best = mAP > best_mAP
best_mAP = max(mAP, best_mAP)
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch + 1,
'best_mAP': mAP,
},
is_best,
fpath=osp.join(args.logs_dir, 'model_best.pth.tar'))
if epoch == args.epochs - 1:
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch + 1,
'best_mAP': mAP,
},
True,
fpath=osp.join(args.logs_dir, 'last.pth.tar'))
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
print(args)
# Create data loaders
if args.big_height is None or args.big_width is None or args.target_height is None or args.target_width is None:
args.big_height, args.big_width, args.target_height, args.target_width = (
256, 256, 224, 224)
dataset, num_classes, train_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.big_height, args.big_width,
args.target_height, args.target_width, args.batch_size, args.num_instances,
args.workers, args.combine_trainval)
# Create models
model = models.create(name=args.arch,
num_classes=num_classes,
num_features=args.features,
norm=True)
print(model)
# Load from checkpoint
start_epoch = best = 0
if args.weights and hasattr(model, 'base'):
print('loading resnet50')
checkpoint = load_checkpoint(args.weights)
del (checkpoint['fc.weight'])
del (checkpoint['fc.bias'])
model.base.load_state_dict(checkpoint)
if args.resume and not args.weights:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
print("=> Start epoch {} best top1 {:.1%}".format(start_epoch, best))
model = nn.DataParallel(model).cuda()
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
print("Test:")
evaluator.evaluate(test_loader, dataset.query, dataset.gallery)
return
# Criterion
if args.arch == 'ResNet50_mgn_lr' or args.arch == 'ResNet101_mgn_lr' or args.arch == 'ResNet152_mgn_lr':
criterion = MGN_loss(margin1=1.2,
num_instances=4,
alpha=1.0,
gamma=1.0,
theta=0.1,
has_trip=True).cuda()
elif args.arch == 'ResNet_reid_50' or args.arch == 'ResNet_reid_101' or args.arch == 'ResNet_reid_152':
# criterion = XentropyLoss_SAC(theta=0.2,gamma=1).cuda()
ranking_loss = nn.MarginRankingLoss(margin=args.margin).cuda()
criterion = { 'XentropyLoss_SAC': XentropyLoss_SAC(theta=0.2,gamma=1).cuda(), \
'trihard': TripletLoss(ranking_loss).cuda() }
else:
criterion = nn.CrossEntropyLoss().cuda()
# Optimizer
frozen_layerName = [
'conv1',
'bn1',
'relu',
'maxpool',
'layer1',
'layer2',
]
##### Optimizer
if args.frozen_sublayer:
frozen_Source = None
if hasattr(model.module, 'base'):
frozen_Source = 'model.module.base.'
elif hasattr(model.module, frozen_layerName[0]):
frozen_Source = 'model.module.'
else:
raise RuntimeError(
'Not freeze layers but frozen_sublayer is True!')
base_params_set = set()
for subLayer in frozen_layerName:
if hasattr(eval(frozen_Source[:-1]), subLayer):
print('frozen layer: ', subLayer)
single_module_param = eval(frozen_Source + subLayer +
'.parameters()')
# base_params.append(single_module_param)
single_module_param_set = set(map(id, single_module_param))
base_params_set = base_params_set | single_module_param_set
else:
print("current model doesn't have ", subLayer)
new_params = [
p for p in model.parameters() if id(p) not in base_params_set
]
base_params = [
p for p in model.parameters() if id(p) in base_params_set
]
param_groups = [{
'params': base_params,
'lr_mult': 0.1
}, {
'params': new_params,
'lr_mult': 1.0
}]
else:
param_groups = model.parameters()
optimizer = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
# Trainer
if args.arch == 'ResNet50_mgn_lr' or args.arch == 'ResNet101_mgn_lr' or args.arch == 'ResNet152_mgn_lr':
trainer = Trainer(model, criterion)
elif args.arch == 'ResNet_reid_50' or args.arch == 'ResNet_reid_101' or args.arch == 'ResNet_reid_152':
trainer = Trainer_SAC_Triplet(
model, criterion, metric_loss_weight=args.metric_loss_weight)
else:
trainer = Cross_Trihard_Trainer(
model, criterion, metric_loss_weight=args.metric_loss_weight)
# Schedule learning rate
print(args.step_epoch)
scheduler = WarmupMultiStepLR(optimizer,
args.step_epoch,
gamma=args.gamma,
warmup_factor=args.warm_up_factor,
warmup_iters=args.warm_up_iter)
# Start training
for epoch in range(start_epoch + 1, args.epochs + 1):
scheduler.step()
trainer.train(epoch, train_loader, optimizer)
if epoch % args.epoch_inter == 0 or epoch >= args.dense_evaluate:
tmp_mAP, tmp_res = evaluator.evaluate(test_loader, dataset.query,
dataset.gallery)
if epoch >= args.start_save:
if tmp_mAP > best:
best = tmp_mAP
flag = True
else:
flag = False
# save_checkpoint({
# 'state_dict': model.module.state_dict(),
# 'epoch': epoch,
# 'best_map':tmp_mAP
# }, flag, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch,
'best_map': tmp_mAP
},
flag,
fpath=osp.join(args.logs_dir, 'pass%d.pth.tar' % (epoch)))
print('Test with best model:')
checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
model.module.load_state_dict(checkpoint['state_dict'])
evaluator.evaluate(test_loader, dataset.query, dataset.gallery)
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch == 'inception' else \
(256, 128)
# get source data
src_dataset, src_extfeat_loader = \
get_source_data(args.src_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# get target data
tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \
get_data(args.tgt_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
# Hacking here to let the classifier be the number of source ids
if args.src_dataset == 'dukemtmc':
model = models.create(args.arch, num_classes=632, pretrained=False)
elif args.src_dataset == 'market1501':
model = models.create(args.arch, num_classes=676, pretrained=False)
else:
raise RuntimeError(
'Please specify the number of classes (ids) of the network.')
# Load from checkpoint
if args.resume:
print(
'Resuming checkpoints from finetuned model on another dataset...\n'
)
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'], strict=False)
else:
raise RuntimeWarning('Not using a pre-trained model.')
model = nn.DataParallel(model).cuda()
# evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery)
# if args.evaluate: return
# Criterion
criterion = [
TripletLoss(args.margin, args.num_instances, isAvg=True,
use_semi=True).cuda(),
TripletLoss(args.margin, args.num_instances, isAvg=True,
use_semi=True).cuda(), None, None
]
# Optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
# training stage transformer on input images
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transformer = T.Compose([
T.Resize((args.height, args.width)),
T.RandomHorizontalFlip(),
T.ToTensor(), normalizer,
T.RandomErasing(probability=0.5, sh=0.2, r1=0.3)
])
evaluator = Evaluator(model, print_freq=args.print_freq)
evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery)
# # Start training
for iter_n in range(args.iteration):
if args.lambda_value == 0:
source_features = 0
else:
# get source datas' feature
source_features, _ = extract_features(model,
src_extfeat_loader,
print_freq=args.print_freq)
# synchronization feature order with src_dataset.train
source_features = torch.cat([
source_features[f].unsqueeze(0)
for f, _, _, _ in src_dataset.train
], 0)
# extract training images' features
print('Iteration {}: Extracting Target Dataset Features...'.format(
iter_n + 1))
target_features, tarNames = extract_features(
model, tgt_extfeat_loader, print_freq=args.print_freq)
# synchronization feature order with dataset.train
target_features = torch.cat([
target_features[f].unsqueeze(0)
for f, _, _, _ in tgt_dataset.trainval
], 0)
# target_real_label = np.asarray([tarNames[f].unsqueeze(0) for f, _, _, _ in tgt_dataset.trainval])
# calculate distance and rerank result
# method 1
target_features = target_features.numpy()
rerank_dist = re_ranking(source_features,
target_features,
lambda_value=args.lambda_value)
# method 2
# distmat_qq = calDis(source_features, source_features)
# distmat_qg = calDis(source_features, target_features)
# distmat_gg = calDis(target_features, target_features)
# rerank_dist = re_ranking2(distmat_qg.numpy(), distmat_qq.numpy(), distmat_gg.numpy())
cluster = HDBSCAN(metric='precomputed', min_samples=10)
# select & cluster images as training set of this epochs
clusterRes = cluster.fit(rerank_dist)
labels, label_num = clusterRes.labels_, clusterRes.labels_.max() + 1
centers = np.zeros((label_num, target_features.shape[1]))
nums = [0] * target_features.shape[1]
print('clusters num =', label_num)
# generate new dataset
new_dataset = []
index = -1
for (fname, _, cam, timestamp), label in zip(tgt_dataset.trainval,
labels):
index += 1
if label == -1: continue
# dont need to change codes in trainer.py _parsing_input function and sampler function after add 0
new_dataset.append((fname, label, cam, timestamp))
centers[label] += target_features[index]
nums[label] += 1
print('Iteration {} have {} training images'.format(
iter_n + 1, len(new_dataset)))
train_loader = DataLoader(Preprocessor(new_dataset,
root=tgt_dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size,
num_workers=4,
sampler=RandomIdentitySampler(
new_dataset, args.num_instances),
pin_memory=True,
drop_last=True)
for i in range(label_num):
centers[i] /= nums[i]
criterion[3] = ClassificationLoss(normalize(centers, axis=1)).cuda()
classOptimizer = torch.optim.Adam(
[{
'params': model.parameters()
}, {
'params': criterion[3].classifier.parameters(),
'lr': 1e-3
}],
lr=args.lr)
class_trainer = ClassificationTrainer(model, train_loader, criterion,
classOptimizer)
for epoch in range(args.epochs):
class_trainer.train(epoch)
rank_score = evaluator.evaluate(test_loader, tgt_dataset.query,
tgt_dataset.gallery)
# Evaluate
rank_score = evaluator.evaluate(test_loader, tgt_dataset.query,
tgt_dataset.gallery)
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch + 1,
'best_top1': rank_score.market1501[0],
},
True,
fpath=osp.join(args.logs_dir, 'adapted.pth.tar'))
return (rank_score.map, rank_score.market1501[0])
def main(args):
# seed
if args.seed is not None:
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
else:
torch.backends.cudnn.benchmark = True
if args.logs_dir is None:
args.logs_dir = osp.join(
f'logs/triplet/{args.dataset}',
datetime.datetime.today().strftime('%Y-%m-%d_%H-%M-%S'))
else:
args.logs_dir = osp.join(f'logs/triplet/{args.dataset}', args.logs_dir)
if args.train:
os.makedirs(args.logs_dir, exist_ok=True)
copy_tree('./reid', args.logs_dir + '/scripts/reid')
for script in os.listdir('.'):
if script.split('.')[-1] == 'py':
dst_file = os.path.join(args.logs_dir, 'scripts',
os.path.basename(script))
shutil.copyfile(script, dst_file)
sys.stdout = Logger(os.path.join(args.logs_dir, 'log.txt'), )
print('Settings:')
print(vars(args))
print('\n')
# Create data loaders
assert args.num_instances > 1, "num_instances should be larger than 1"
assert args.batch_size % args.num_instances == 0, 'num_instances should divide batch_size'
dataset, num_classes, train_loader, query_loader, gallery_loader, _ = \
get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.num_workers,
args.combine_trainval, args.crop, args.tracking_icams, args.tracking_fps, args.re, args.num_instances,
False)
# Create model for triplet (num_classes = 0, num_instances > 0)
model = models.create('ide',
feature_dim=args.feature_dim,
num_classes=0,
norm=args.norm,
dropout=args.dropout,
last_stride=args.last_stride)
# Load from checkpoint
start_epoch = best_top1 = 0
if args.resume:
resume_fname = osp.join(f'logs/triplet/{args.dataset}', args.resume,
'model_best.pth.tar')
model, start_epoch, best_top1 = checkpoint_loader(model, resume_fname)
print("=> Last epoch {} best top1 {:.1%}".format(
start_epoch, best_top1))
start_epoch += 1
model = nn.DataParallel(model).cuda()
# Criterion
criterion = TripletLoss(margin=args.margin).cuda()
# Optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# Trainer
trainer = Trainer(model, criterion)
# Evaluator
evaluator = Evaluator(model)
if args.train:
# Schedule learning rate
def adjust_lr(epoch):
if epoch <= args.step_size:
lr = args.lr
else:
lr = args.lr * (0.001**(float(epoch - args.step_size) /
(args.epochs - args.step_size)))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
# Draw Curve
epoch_s = []
loss_s = []
prec_s = []
eval_epoch_s = []
eval_top1_s = []
# Start training
for epoch in range(start_epoch + 1, args.epochs + 1):
adjust_lr(epoch)
# train_loss, train_prec = 0, 0
train_loss, train_prec = trainer.train(epoch,
train_loader,
optimizer,
fix_bn=args.fix_bn)
if epoch < args.start_save:
continue
if epoch % 25 == 0:
top1 = evaluator.evaluate(query_loader, gallery_loader,
dataset.query, dataset.gallery)
eval_epoch_s.append(epoch)
eval_top1_s.append(top1)
else:
top1 = 0
is_best = top1 >= best_top1
best_top1 = max(top1, best_top1)
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch,
'best_top1': best_top1,
},
is_best,
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
epoch_s.append(epoch)
loss_s.append(train_loss)
prec_s.append(train_prec)
draw_curve(os.path.join(args.logs_dir, 'train_curve.jpg'), epoch_s,
loss_s, prec_s, eval_epoch_s, None, eval_top1_s)
pass
# Final test
print('Test with best model:')
model, start_epoch, best_top1 = checkpoint_loader(
model, osp.join(args.logs_dir, 'model_best.pth.tar'))
print("=> Start epoch {} best top1 {:.1%}".format(
start_epoch, best_top1))
evaluator.evaluate(query_loader, gallery_loader, dataset.query,
dataset.gallery)
else:
print("Test:")
evaluator.evaluate(query_loader, gallery_loader, dataset.query,
dataset.gallery)
pass
from __future__ import print_function, absolute_import
import time
import torch
from reid.evaluator import accuracy
from utils.meters import AverageMeter
import torch.nn.functional as F
from tensorboardX import SummaryWriter
from visualize import reverse_normalize
from cam_functions import visual_batch
from reid.loss import TripletLoss, TripletLoss_OIM
criterion_triplet_oim = TripletLoss_OIM('soft', True)
criterion_triplet = TripletLoss('soft', True)
class BaseTrainer(object):
def __init__(self, model, criterion):
super(BaseTrainer, self).__init__()
self.model = model
self.criterion_ver = criterion
self.criterion_ver_uncorr = criterion
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
def train(self, epoch, data_loader, optimizer1):
self.model.train()
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
def train(self,
train_data,
step,
epochs=70,
step_size=55,
init_lr=0.1,
dropout=0.5):
""" create model and dataloader """
if self.model:
model = self.model
else:
model = models.create(self.model_name,
dropout=self.dropout,
num_classes=self.num_classes,
mode=self.mode)
model = nn.DataParallel(model).cuda()
dataloader = self.get_dataloader(train_data, training=True)
# the base parameters for the backbone (e.g. ResNet50)
base_param_ids = set(map(id, model.module.base.parameters()))
# we fixed the first three blocks to save GPU memory
# base_params_need_for_grad = filter(lambda p: p.requires_grad, model.module.base.parameters())
# params of the new layers
new_params = [
p for p in model.parameters() if id(p) not in base_param_ids
]
# set the learning rate for backbone to be 0.1 times
param_groups = [
{
'params': model.module.base.parameters(),
'lr_mult': 1.0
},
# {'params': base_params_need_for_grad, 'lr_mult': 0.1},
{
'params': new_params,
'lr_mult': 1.0
}
]
criterion = []
if self.num_classes == 0:
criterion.append(
TripletLoss(margin=0.3,
num_instances=self.num_instances).cuda())
criterion.append(
TripletLoss(margin=0.3,
num_instances=self.num_instances).cuda())
trainer = Trainer(model, criterion)
else:
criterion.append(nn.CrossEntropyLoss().cuda())
criterion.append(nn.MSELoss().cuda())
trainer = DistillTrainer(model, self.model_distill, criterion)
optimizer = torch.optim.SGD(param_groups,
lr=init_lr,
momentum=0.9,
weight_decay=5e-4,
nesterov=True)
# change the learning rate by step
def adjust_lr(epoch, step_size):
lr = init_lr / (10**(epoch // step_size))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
if epoch % step_size == 0:
print("Epoch {}, current lr {}".format(epoch, lr))
# def adjust_lr(epoch):
# if epoch <=7:
# lr = args.lr
# elif epoch <= 14:
# lr = 0.3 * args.lr
# else:
# lr = 0.1 * args.lr
# for g in optimizer.param_groups:
# g['lr'] = lr * g.get('lr_mult', 1)
""" main training process """
for epoch in range(epochs):
adjust_lr(epoch, step_size)
trainer.train(epoch, dataloader, optimizer, print_freq=20)
self.model = model
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 main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch == 'inception' else \
(256, 128)
# get source data
src_dataset, src_extfeat_loader = \
get_source_data(args.src_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# get target data
tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \
get_data(args.tgt_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
# Hacking here to let the classifier be the number of source ids
if args.src_dataset == 'dukemtmc':
model = models.create(args.arch, num_classes=632, pretrained=False)
elif args.src_dataset == 'market1501':
model = models.create(args.arch, num_classes=676, pretrained=False)
else:
raise RuntimeError(
'Please specify the number of classes (ids) of the network.')
# Load from checkpoint
if args.resume:
print(
'Resuming checkpoints from finetuned model on another dataset...\n'
)
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'], strict=False)
else:
raise RuntimeWarning('Not using a pre-trained model.')
model = nn.DataParallel(model).cuda()
# Distance metric
# metric = DistanceMetric(algorithm=args.dist_metric)
# Evaluator
evaluator = Evaluator(model, print_freq=args.print_freq)
print(
"Test with the original model trained on source domain (direct transfer):"
)
rank_score_best = evaluator.evaluate(test_loader, tgt_dataset.query,
tgt_dataset.gallery)
best_map = rank_score_best.map #market1501[0]-->rank-1
if args.evaluate:
return
# Criterion
criterion = [
TripletLoss(args.margin, args.num_instances).cuda(),
TripletLoss(args.margin, args.num_instances).cuda(),
AccumulatedLoss(args.margin, args.num_instances).cuda(),
nn.CrossEntropyLoss().cuda()
]
# Optimizer
optimizer = torch.optim.SGD(
model.parameters(),
lr=args.lr,
momentum=0.9,
)
# training stage transformer on input images
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transformer = T.Compose([
T.Resize((args.height, args.width)),
T.RandomHorizontalFlip(),
T.ToTensor(), normalizer,
T.RandomErasing(probability=0.5, sh=0.2, r1=0.3)
])
# Start training
for iter_n in range(args.iteration):
if args.lambda_value == 0:
source_features = 0 #this value controls the usage of source data
else:
# get source datas' feature
source_features, _ = extract_features(model,
src_extfeat_loader,
print_freq=args.print_freq)
# synchronization feature order with src_dataset.train
source_features = torch.cat([
source_features[f].unsqueeze(0)
for f, _, _ in src_dataset.train
], 0)
# extract training images' features
print('Iteration {}: Extracting Target Dataset Features...'.format(
iter_n + 1))
target_features, _ = extract_features(model,
tgt_extfeat_loader,
print_freq=args.print_freq)
# synchronization feature order with dataset.train
target_features = torch.cat([
target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval
], 0)
# calculate distance and rerank result
print('Calculating feature distances...')
target_features = target_features.numpy()
rerank_dist = re_ranking(source_features,
target_features,
lambda_value=args.lambda_value)
if iter_n == 0:
# DBSCAN cluster
tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2
tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1
tri_mat = np.sort(tri_mat, axis=None)
top_num = np.round(args.rho * tri_mat.size).astype(int)
eps = tri_mat[:top_num].mean()
print('eps in cluster: {:.3f}'.format(eps))
cluster = DBSCAN(eps=eps,
min_samples=4,
metric='precomputed',
n_jobs=8)
# HDBSCAN cluster
import hdbscan
cluster_hdbscan = hdbscan.HDBSCAN(min_cluster_size=10,
min_samples=4,
metric='precomputed')
# select & cluster images as training set of this epochs
print('Clustering and labeling...')
if args.use_hdbscan_clustering:
print(
'Use the better chlustering algorithm HDBSCAN for clustering'
)
labels = cluster_hdbscan.fit_predict(rerank_dist)
else:
print('Use DBSCAN for clustering')
labels = cluster.fit_predict(rerank_dist)
num_ids = len(set(labels)) - 1
print('Only do once, Iteration {} have {} training ids'.format(
iter_n + 1, num_ids))
# generate new dataset
new_dataset = []
for (fname, _, _), label in zip(tgt_dataset.trainval, labels):
if label == -1:
continue
# dont need to change codes in trainer.py _parsing_input function and sampler function after add 0
new_dataset.append((fname, label, 0))
print('Only do once, Iteration {} have {} training images'.format(
iter_n + 1, len(new_dataset)))
train_loader = DataLoader(
Preprocessor_return_index(new_dataset,
root=tgt_dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size,
num_workers=4,
sampler=RandomIdentitySampler(new_dataset, args.num_instances),
pin_memory=True,
drop_last=True)
# init pseudo/fake labels, y_tilde in cvpr19's paper:
new_label = np.zeros([len(new_dataset), num_ids])
# init y_tilde, let softmax(y_tilde) is noisy labels
for index, (imgs, _, pids, _, index) in enumerate(train_loader):
index = index.numpy()
onehot = torch.zeros(pids.size(0),
num_ids).scatter_(1, pids.view(-1, 1),
10.0)
onehot = onehot.numpy()
new_label[index, :] = onehot
# Using clustered label to init the new classifier:
classifier = nn.Linear(2048, num_ids, bias=False)
classifier.apply(weights_init_classifier)
classifier = nn.DataParallel(classifier).cuda()
optimizer_cla = torch.optim.SGD(classifier.parameters(),
lr=args.lr * 10,
momentum=0.9)
# train model with new generated dataset
trainer = Trainer_with_learnable_label(model,
classifier,
criterion,
print_freq=args.print_freq)
evaluator = Evaluator(model, print_freq=args.print_freq)
# Start training
for epoch in range(args.epochs):
trainer.train(epoch, train_loader, new_label, optimizer,
optimizer_cla)
# Evaluate
rank_score = evaluator.evaluate(test_loader, tgt_dataset.query,
tgt_dataset.gallery)
#Save the best ckpt:
rank1 = rank_score.market1501[0]
mAP = rank_score.map
is_best_mAP = mAP > best_map
best_map = max(mAP, best_map)
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': iter_n + 1,
'best_mAP': best_map,
# 'num_ids': num_ids,
},
is_best_mAP,
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print(
'\n * Finished epoch {:3d} top1: {:5.1%} mAP: {:5.1%} best_mAP: {:5.1%}{}\n'
.format(iter_n + 1, rank1, mAP, best_map,
' *' if is_best_mAP else ''))
return (rank_score.map, rank_score.market1501[0])
def main(args):
np.random.seed(args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.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'))
# Log args
print(args)
args_dict = vars(args)
with open(osp.join(args.logs_dir, 'args.json'), 'w') as f:
json.dump(args_dict, f)
# Create data loaders
dataset, num_classes, train_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.height, \
args.width, args.crop_height, args.crop_width, args.batch_size, \
args.caffe_sampler, \
args.workers)
# Create model
valid_args = ['features', 'use_relu', 'dilation']
model_kwargs = {k:v for k,v in args_dict.items() if k in valid_args}
model = models.create(args.arch, **model_kwargs)
model = nn.DataParallel(model).cuda()
# Evaluator
evaluator = Evaluator(model)
# Criterion
criterion = TripletLoss(margin=args.margin).cuda()
# Optimizer
params = []
params_dict = dict(model.named_parameters())
for key, value in params_dict.items():
if value.requires_grad == False:
continue
if key[-4:]=='bias':
params += [{'params': [value], 'lr':args.lr*2, 'weight_decay':args.weight_decay*0.0}]
else:
params += [{'params': [value], 'lr':args.lr*1, 'weight_decay':args.weight_decay*1.0}]
optimizer = SGD_caffe(params,
lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum)
# Trainer
trainer = Trainer(model, criterion)
# Evaluate
def evaluate(test_model):
print('Test with model {}:'.format(test_model))
checkpoint = load_checkpoint(osp.join(args.logs_dir, '{}.pth.tar'.format(test_model)))
model.module.load(checkpoint)
evaluator.evaluate(test_loader, dataset.query, dataset.gallery, msg='TEST')
# Schedule learning rate
def adjust_lr(epoch):
lr = args.lr if epoch <= 200 else \
args.lr * (0.2 ** ((epoch-200)//200 + 1))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
# Start training
for epoch in range(1, args.epochs+1):
adjust_lr(epoch)
trainer.train(epoch, train_loader, optimizer)
# Save
if epoch%200 == 0 or epoch==args.epochs:
save_dict = model.module.save_dict()
save_dict.update({'epoch': epoch})
save_checkpoint(save_dict, fpath=osp.join(args.logs_dir, 'epoch_{}.pth.tar'.format(epoch)))
print('\n * Finished epoch {:3d}\n'.format(epoch))
evaluate('epoch_750')
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
else:
log_dir = osp.dirname(args.resume)
sys.stdout = Logger(osp.join(log_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
# Create data loaders with specific size of photos
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch in \
['inception', 'inceptionNet', 'squeezenet', 'squeezenet1_0', 'squeezenet1_1'] else \
(384, 128) if args.arch == 'mgn' else (160, 64) if args.arch == 'hacnn' else\
(256, 128)
dataset, train_loader, val_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.height,
args.width, args.batch_size, args.workers,
args.combine_trainval, args.np_ratio)
# Create model
args.features = 1024 if args.arch in \
['squeezenet', 'squeezenet1_0', 'squeezenet1_1'] else \
1536 if args.arch in ['mgn', 'inception', 'inceptionv4'] else \
2432 if args.arch == 'hacnn' else \
2048
base_model = models.create(args.arch, cut_at_pooling=True)
embed_model = EltwiseSubEmbed(use_batch_norm=True,
use_classifier=True,
num_features=args.features,
num_classes=2,
use_sft=args.spectral)
model = SiameseNet(base_model, embed_model)
model = nn.DataParallel(model).cuda() # gpu #
print(model)
print('No of parameters:',
sum(p.numel() for p in model.parameters() if p.requires_grad))
# Evaluator
evaluator = CascadeEvaluator(
torch.nn.DataParallel(base_model).cuda(), # gpu #
embed_model,
embed_dist_fn=lambda x: F.softmax(Variable(x), dim=1).data[:, 0])
# Load from checkpoint
best_mAP = 0
if args.resume:
checkpoint = load_checkpoint(args.resume)
if 'state_dict' in checkpoint.keys():
checkpoint = checkpoint['state_dict']
model.load_state_dict(checkpoint)
print("Test the loaded model:")
top1, mAP = evaluator.evaluate(test_loader,
dataset.query,
dataset.gallery,
rerank_topk=100,
dataset=args.dataset)
best_mAP = mAP
if args.evaluate:
return
# Criterion
if args.criterion == 'cross':
criterion = nn.CrossEntropyLoss().cuda()
elif args.criterion == 'triplet':
criterion = TripletLoss(args.margin).cuda()
# Optimizer
param_groups = [{
'params': model.module.base_model.parameters(),
'lr_mult': 1.0
}, {
'params': model.module.embed_model.parameters(),
'lr_mult': 1.0
}]
optimizer = torch.optim.SGD(param_groups,
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Trainer
trainer = SiameseTrainer(model, criterion)
# Schedule learning rate
def adjust_lr(epoch):
lr = args.lr * (0.1**(epoch // args.step_size))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
# Start training
for epoch in range(0, args.epochs):
adjust_lr(epoch)
trainer.train(epoch, train_loader, optimizer, base_lr=args.lr)
if epoch % args.eval_step == 0:
mAP = evaluator.evaluate(val_loader,
dataset.val,
dataset.val,
top1=False)
is_best = mAP > best_mAP
best_mAP = max(mAP, best_mAP)
save_checkpoint({'state_dict': model.state_dict()},
is_best,
fpath=osp.join(args.logs_dir,
'checkpoint.pth.tar'))
print('\n * Finished epoch {:3d} mAP: {:5.1%} best: {:5.1%}{}\n'.
format(epoch, mAP, best_mAP, ' *' 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'])
evaluator.evaluate(test_loader,
dataset.query,
dataset.gallery,
dataset=args.dataset)
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch == 'inception' else \
(256, 128)
# get source data
src_dataset, src_extfeat_loader = \
get_source_data(args.src_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# get target data
tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \
get_data(args.tgt_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
# Hacking here to let the classifier be the number of source ids
if args.src_dataset == 'dukemtmc':
model = models.create(args.arch, num_classes=632, pretrained=False)
elif args.src_dataset == 'market1501':
model = models.create(args.arch, num_classes=676, pretrained=False)
else:
raise RuntimeError(
'Please specify the number of classes (ids) of the network.')
# Load from checkpoint
if args.resume:
print(
'Resuming checkpoints from finetuned model on another dataset...\n'
)
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'], strict=False)
else:
raise RuntimeWarning('Not using a pre-trained model.')
model = nn.DataParallel(model).cuda()
# Criterion
criterion = [
TripletLoss(args.margin, args.num_instances, use_semi=False).cuda(),
TripletLoss(args.margin, args.num_instances, use_semi=False).cuda()
]
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
# training stage transformer on input images
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transformer = T.Compose([
T.Resize((args.height, args.width)),
T.RandomHorizontalFlip(),
T.ToTensor(), normalizer,
T.RandomErasing(probability=0.5, sh=0.2, r1=0.3)
])
# # Start training
for iter_n in range(args.iteration):
if args.lambda_value == 0:
source_features = 0
else:
# get source datas' feature
source_features, _ = extract_features(model,
src_extfeat_loader,
print_freq=args.print_freq)
# synchronization feature order with src_dataset.train
source_features = torch.cat([
source_features[f].unsqueeze(0)
for f, _, _ in src_dataset.train
], 0)
# extract training images' features
print('Iteration {}: Extracting Target Dataset Features...'.format(
iter_n + 1))
target_features, _ = extract_features(model,
tgt_extfeat_loader,
print_freq=args.print_freq)
# synchronization feature order with dataset.train
target_features = torch.cat([
target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval
], 0)
# calculate distance and rerank result
print('Calculating feature distances...')
target_features = target_features.numpy()
rerank_dist = re_ranking(source_features,
target_features,
lambda_value=args.lambda_value)
if iter_n == 0:
# DBSCAN cluster
tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2
tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1
tri_mat = np.sort(tri_mat, axis=None)
top_num = np.round(args.rho * tri_mat.size).astype(int)
eps = tri_mat[:top_num].mean()
print('eps in cluster: {:.3f}'.format(eps))
cluster = DBSCAN(eps=eps,
min_samples=4,
metric='precomputed',
n_jobs=8)
# select & cluster images as training set of this epochs
print('Clustering and labeling...')
labels = cluster.fit_predict(rerank_dist)
num_ids = len(set(labels)) - 1
print('Iteration {} have {} training ids'.format(iter_n + 1, num_ids))
# generate new dataset
new_dataset = []
# assign label for target ones
newLab = labelNoise(torch.from_numpy(target_features),
torch.from_numpy(labels))
# unknownFeats = target_features[labels==-1,:]
counter = 0
from collections import defaultdict
realIDs, fakeIDs = defaultdict(list), []
for (fname, realID, cam), label in zip(tgt_dataset.trainval, newLab):
# dont need to change codes in trainer.py _parsing_input function and sampler function after add 0
new_dataset.append((fname, label, cam))
realIDs[realID].append(counter)
fakeIDs.append(label)
counter += 1
precision, recall, fscore = calScores(realIDs, np.asarray(fakeIDs))
print('Iteration {} have {} training images'.format(
iter_n + 1, len(new_dataset)))
print(
f'precision:{precision * 100}, recall:{100 * recall}, fscore:{fscore}'
)
train_loader = DataLoader(Preprocessor(new_dataset,
root=tgt_dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size,
num_workers=4,
sampler=RandomIdentitySampler(
new_dataset, args.num_instances),
pin_memory=True,
drop_last=True)
trainer = Trainer(model, criterion)
# Start training
for epoch in range(args.epochs):
trainer.train(epoch, train_loader, optimizer) # to at most 80%
# test only
evaluator = Evaluator(model, print_freq=args.print_freq)
# rank_score = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery)
# Evaluate
rank_score = evaluator.evaluate(test_loader, tgt_dataset.query,
tgt_dataset.gallery)
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch + 1,
'best_top1': rank_score.market1501[0],
},
True,
fpath=osp.join(args.logs_dir, 'adapted.pth.tar'))
return rank_score.map, rank_score.market1501[0]
def main(args):
np.random.seed(args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.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'))
# Log args
print(args)
args_dict = vars(args)
with open(osp.join(args.logs_dir, 'args.json'), 'w') as f:
json.dump(args_dict, f)
# Create data loaders
dataset, num_classes, train_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.height, \
args.width, args.crop_height, args.crop_width, args.batch_size, \
args.caffe_sampler, \
args.workers)
# Create model
net_base = SSTTracker()
net_finetune = Finetune().cuda()
net_finetune.init_weights()
# Criterion
criterion = TripletLoss(margin=args.margin).cuda()
# Optimizer
optimizer = SGD_caffe(net_finetune.parameters(),
lr=args.lr,
weight_decay=args.weight_decay,
momentum=args.momentum)
# Schedule learning rate
def adjust_lr(epoch):
lr = args.lr * (0.1**(epoch // 20))
for g in optimizer.param_groups:
g['lr'] = lr
# lr = args.lr if epoch <= 200 else \
# args.lr * (0.2 ** ((epoch-200)//200 + 1))
# for g in optimizer.param_groups:
# g['lr'] = lr * g.get('lr_mult', 1)
def parse_data(inputs):
imgs, _, pids, _ = inputs
inputs = Variable(imgs).cuda()
targets = Variable(pids.cuda())
return inputs, targets
for epoch in range(750):
adjust_lr(epoch)
loss_epoch_train = 0.0
net_finetune.train()
for i, inputs in enumerate(tqdm(train_loader)):
optimizer.zero_grad()
inputs, targets = parse_data(inputs)
outputs = inputs
for k in range(35):
outputs = net_base.sst.vgg[k](outputs)
outputs = net_finetune(outputs).squeeze()
loss, prec, num_eff, num_all = criterion(outputs, targets)
loss_epoch_train += loss.data
loss.backward()
optimizer.step()
loss_epoch_train /= 41
loss_epoch_valid = 0.0
net_finetune.eval()
for i, inputs in enumerate(tqdm(test_loader)):
inputs, targets = parse_data(inputs)
outputs = inputs
for k in range(35):
outputs = net_base.sst.vgg[k](outputs)
outputs = net_finetune(outputs).squeeze()
loss, prec, num_eff, num_all = criterion(outputs, targets)
loss_epoch_valid += loss.data
loss_epoch_valid /= 108
print('epoch:%d__train_loss:%.4f__valid_loss:%.4f' %
(epoch, loss_epoch_train.data, loss_epoch_valid.data))
torch.save(net_finetune.cpu().state_dict(),
"./weights/reid_finetune_{}.pth".format(epoch + 1))
net_finetune.cuda()
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch == 'inception' else \
(256, 128)
# get source data
src_dataset, src_extfeat_loader = \
get_source_data(args.src_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# get target data
tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \
get_data(args.tgt_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
# Hacking here to let the classifier be the number of source ids
if args.src_dataset == 'dukemtmc':
model = models.create(args.arch, num_classes=632, pretrained=False)
elif args.src_dataset == 'market1501':
model = models.create(args.arch, num_classes=676, pretrained=False)
else:
raise RuntimeError(
'Please specify the number of classes (ids) of the network.')
# Load from checkpoint
if args.resume:
print(
'Resuming checkpoints from finetuned model on another dataset...\n'
)
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'], strict=False)
else:
raise RuntimeWarning('Not using a pre-trained model.')
model = nn.DataParallel(model).cuda()
# Distance metric
metric = DistanceMetric(algorithm=args.dist_metric)
# Evaluator
evaluator = Evaluator(model, print_freq=args.print_freq)
print(
"Test with the original model trained on target domain (direct transfer):"
)
evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery)
if args.evaluate:
return
# Criterion
criterion = [
TripletLoss(args.margin, args.num_instances).cuda(),
TripletLoss(args.margin, args.num_instances).cuda(),
]
# Optimizer
optimizer = torch.optim.SGD(
model.parameters(),
lr=args.lr,
momentum=0.9,
)
# training stage transformer on input images
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transformer = T.Compose([
T.Resize((args.height, args.width)),
T.RandomHorizontalFlip(),
T.ToTensor(), normalizer,
T.RandomErasing(probability=0.5, sh=0.2, r1=0.3)
])
# Start training
for iter_n in range(args.iteration):
if args.lambda_value == 0:
source_features = 0
else:
# get source datas' feature
source_features, _ = extract_features(model,
src_extfeat_loader,
print_freq=args.print_freq)
# synchronization feature order with src_dataset.train
source_features = torch.cat([
source_features[f].unsqueeze(0)
for f, _, _ in src_dataset.train
], 0)
# extract training images' features
print('Iteration {}: Extracting Target Dataset Features...'.format(
iter_n + 1))
target_features, _ = extract_features(model,
tgt_extfeat_loader,
print_freq=args.print_freq)
# synchronization feature order with dataset.train
target_features = torch.cat([
target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval
], 0)
# calculate distance and rerank result
print('Calculating feature distances...')
target_features = target_features.numpy()
rerank_dist = re_ranking(source_features,
target_features,
lambda_value=args.lambda_value)
if iter_n == 0:
# DBSCAN cluster
tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2
tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1
tri_mat = np.sort(tri_mat, axis=None)
top_num = np.round(args.rho * tri_mat.size).astype(int)
eps = tri_mat[:top_num].mean()
print('eps in cluster: {:.3f}'.format(eps))
cluster = DBSCAN(eps=eps,
min_samples=4,
metric='precomputed',
n_jobs=8)
# select & cluster images as training set of this epochs
print('Clustering and labeling...')
labels = cluster.fit_predict(rerank_dist)
num_ids = len(set(labels)) - 1
print('Iteration {} have {} training ids'.format(iter_n + 1, num_ids))
# generate new dataset
new_dataset = []
for (fname, _, _), label in zip(tgt_dataset.trainval, labels):
if label == -1:
continue
# dont need to change codes in trainer.py _parsing_input function and sampler function after add 0
new_dataset.append((fname, label, 0))
print('Iteration {} have {} training images'.format(
iter_n + 1, len(new_dataset)))
train_loader = DataLoader(Preprocessor(new_dataset,
root=tgt_dataset.images_dir,
transform=train_transformer),
batch_size=args.batch_size,
num_workers=4,
sampler=RandomIdentitySampler(
new_dataset, args.num_instances),
pin_memory=True,
drop_last=True)
# train model with new generated dataset
trainer = Trainer(model, criterion, print_freq=args.print_freq)
evaluator = Evaluator(model, print_freq=args.print_freq)
# Start training
for epoch in range(args.epochs):
trainer.train(epoch, train_loader, optimizer)
# Evaluate
rank_score = evaluator.evaluate(test_loader, tgt_dataset.query,
tgt_dataset.gallery)
return (rank_score.map, rank_score.market1501[0])
def main(args):
fix(args.seed)
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
print(args)
# Create data loaders
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 = \
get_data(args.data_dir, args.height,
args.width, args.batch_size, args.num_instance, args.re, args.workers)
# Create model
Encoder, Transfer, CamDis = models.create(args.arch, num_features=args.features,
dropout=args.dropout, num_classes=num_classes)
invNet = InvNet(args.features, num_classes, args.batch_size, beta=args.beta, knn=args.knn, alpha=args.alpha).cuda()
# Load from checkpoint
start_epoch = 0
if args.resume:
checkpoint = load_checkpoint(args.resume)
Encoder.load_state_dict(checkpoint['Encoder'])
Transfer.load_state_dict(checkpoint['Transfer'])
CamDis.load_state_dict(checkpoint['CamDis'])
invNet.load_state_dict(checkpoint['InvNet'])
start_epoch = checkpoint['epoch']
Encoder = Encoder.cuda()
Transfer = Transfer.cuda()
CamDis = CamDis.cuda()
model = [Encoder, Transfer, CamDis]
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
# -----------------------------
v = evaluator.eval_viper(viper_query_loader, viper_gallery_loader, test_dataset.viper_query, test_dataset.viper_gallery, args.output_feature, seed=97)
p = evaluator.eval_prid(prid_query_loader, prid_gallery_loader, test_dataset.prid_query, test_dataset.prid_gallery, args.output_feature, seed=40)
g = evaluator.eval_grid(grid_query_loader, grid_gallery_loader, test_dataset.grid_query, test_dataset.grid_gallery, args.output_feature, seed=28)
l = evaluator.eval_ilids(ilid_query_loader, test_dataset.ilid_query, args.output_feature, seed=24)
# -----------------------------
criterion = []
criterion.append(nn.CrossEntropyLoss().cuda())
criterion.append(TripletLoss(margin=args.margin))
# Optimizer
base_param_ids = set(map(id, Encoder.base.parameters()))
new_params = [p for p in Encoder.parameters() if
id(p) not in base_param_ids]
param_groups = [
{'params': Encoder.base.parameters(), 'lr_mult': 0.1},
{'params': new_params, 'lr_mult': 1.0}]
optimizer_Encoder = torch.optim.SGD(param_groups, lr=args.lr,
momentum=0.9, weight_decay=5e-4, nesterov=True)
# ====
base_param_ids = set(map(id, Transfer.base.parameters()))
new_params = [p for p in Transfer.parameters() if
id(p) not in base_param_ids]
param_groups = [
{'params': Transfer.base.parameters(), 'lr_mult': 0.1},
{'params': new_params, 'lr_mult': 1.0}]
optimizer_Transfer = torch.optim.SGD(param_groups, lr=args.lr,
momentum=0.9, weight_decay=5e-4, nesterov=True)
# ====
param_groups = [
{'params':CamDis.parameters(), 'lr_mult':1.0},
]
optimizer_Cam = torch.optim.SGD(param_groups, lr=args.lr,momentum=0.9, weight_decay=5e-4, nesterov=True)
optimizer = [optimizer_Encoder, optimizer_Transfer, optimizer_Cam]
# Trainer
trainer = Trainer(model, criterion, InvNet=invNet)
# Schedule learning rate
def adjust_lr(epoch):
step_size = 40
lr = args.lr * (0.1 ** ((epoch) // step_size))
for g in optimizer_Encoder.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
for g in optimizer_Transfer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
for g in optimizer_Cam.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
# Start training
for epoch in range(start_epoch, args.epochs):
adjust_lr(epoch)
trainer.train(epoch, source_train_loader, optimizer, args.tri_weight, args.adv_weight, args.mem_weight)
save_checkpoint({
'Encoder': Encoder.state_dict(),
'Transfer': Transfer.state_dict(),
'CamDis': CamDis.state_dict(),
'InvNet': invNet.state_dict(),
'epoch': epoch + 1,
}, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
evaluator = Evaluator(model)
print('\n * Finished epoch {:3d} \n'.
format(epoch))
# Final test
print('Test with best model:')
evaluator = Evaluator(model)
evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, args.output_feature, args.rerank)
[recoloradv_threat, additive_threat],
ap.PerturbationParameters(norm_weights=[1.0, 0.0]),
)
if args.noise_type==0:
noise = combined_threat()
elif args.noise_type==1:
noise = recoloradv_threat()
elif args.noise_type == 2:
noise = additive_threat()
####criterion
criterion = [
TripletLoss(0.5, args.num_instances, False).cuda(),
nn.CrossEntropyLoss().cuda()
]
# multi lr
base_param_ids = set(map(id, model.module.base.params()))
new_params = [p for p in model.module.params() if id(p) not in base_param_ids]
param_groups = [
{'params': model.module.base.params(), 'lr_mult': 1.0},
{'params': new_params, 'lr_mult': 3.0}]
# Optimizer
optimizer_defense = torch.optim.Adam(param_groups, lr=args.lr, weight_decay=args.weight_decay)
# Schedule learning rate
def adjust_lr(epoch):
lr = args.lr if epoch <= 40 else \
args.lr * (0.1 ** ((epoch - 40) / 40.0))
def main(args):
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
dataset, num_classes, source_train_loader, source_triplet_train_loader, \
target_train_loader, query_loader, gallery_loader = \
get_data(args.data_dir, args.source, args.target, args.height,
args.width, args.batch_size, args.triplet_batch_size, args.num_instances, args.target_batch_size, args.re, args.workers)
# Create model
model = models.create(args.arch,
num_features=args.features,
dropout=args.dropout,
num_classes=num_classes,
triplet_features=args.triplet_features)
# Load from checkpoint
start_epoch = 0
if args.resume:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
print("=> Start epoch {} ".format(start_epoch))
# model = nn.DataParallel(model).cuda()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model.to(device)
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
print("Test:")
evaluator.evaluate(query_loader, gallery_loader, dataset.query,
dataset.gallery, args.output_feature, args.rerank)
return
# Criterion
# cross-entropy loss
criterion_c = nn.CrossEntropyLoss().to(device)
# triplet loss
criterion_t = TripletLoss(margin=args.margin).to(device)
# Optimizer
if hasattr(model.module, 'base'):
base_param_ids = set(map(id, model.module.base.parameters())) \
| set(map(id, model.module.triplet.parameters())) \
| set(map(id, model.module.feat.parameters())) \
| set(map(id, model.module.feat_bn.parameters()))
new_params = [
p for p in model.parameters() if id(p) not in base_param_ids
]
param_groups = [{
'params': model.module.base.parameters(),
'lr_mult': 0.1
}, {
'params': new_params,
'lr_mult': 1.0
}]
else:
param_groups = model.parameters()
optimizer = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
# Trainer
trainer = HHLTrainer(model, criterion_c, criterion_t, args.beta)
# Schedule learning rate
def adjust_lr(epoch):
step_size = 40
lr = args.lr * (0.1**(epoch // step_size))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
# Start training
for epoch in range(start_epoch, args.epochs):
adjust_lr(epoch)
trainer.train(epoch, source_train_loader, source_triplet_train_loader,
target_train_loader, optimizer)
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch + 1,
},
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print('\n * Finished epoch {:3d} \n'.format(epoch))
# Final test
print('Test with best model:')
evaluator = Evaluator(model)
evaluator.evaluate(query_loader, gallery_loader, dataset.query,
dataset.gallery, args.output_feature, args.rerank)
def main(args):
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch == 'inception' else \
(256, 128)
## get_source_data
src_dataset, src_extfeat_loader = \
get_source_data(args.src_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# get_target_data
tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \
get_data(args.tgt_dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
# Hacking here to let the classifier be the last feature embedding layer
# Net structure: avgpool -> FC(2048) -> FC(args.features)
num_class = 0
if args.src_dataset == 'dukemtmc':
model = models.create(args.arch, num_classes=num_class, num_split=args.num_split, cluster=args.dce_loss) #duke
elif args.src_dataset == 'market1501':
model = models.create(args.arch, num_classes=num_class, num_split=args.num_split, cluster=args.dce_loss)
else:
raise RuntimeError('Please specify the number of classes (ids) of the network.')
# Load from checkpoint
start_epoch = best_top1 = 0
if args.resume:
print('Resuming checkpoints from finetuned model on another dataset...\n')
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint, strict=False)
else:
raise RuntimeWarning('Not using a pre-trained model')
model = nn.DataParallel(model).cuda()
# Distance metric
metric = DistanceMetric(algorithm=args.dist_metric)
# Evaluator
evaluator = Evaluator(model, print_freq=args.print_freq)
print("Test with the original model trained on source domain:")
best_top1 = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery)
if args.evaluate:
return
# Criterion
criterion = []
criterion.append(TripletLoss(margin=args.margin,num_instances=args.num_instances).cuda())
criterion.append(TripletLoss(margin=args.margin,num_instances=args.num_instances).cuda())
#multi lr
base_param_ids = set(map(id, model.module.base.parameters()))
new_params = [p for p in model.parameters() if
id(p) not in base_param_ids]
param_groups = [
{'params': model.module.base.parameters(), 'lr_mult': 1.0},
{'params': new_params, 'lr_mult': 1.0}]
# Optimizer
optimizer = torch.optim.SGD(param_groups, lr=args.lr,
momentum=0.9, weight_decay=args.weight_decay)
##### adjust lr
def adjust_lr(epoch):
if epoch <= 7:
lr = args.lr
elif epoch <=14:
lr = 0.3 * args.lr
else:
lr = 0.1 * args.lr
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
##### training stage transformer on input images
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transformer = T.Compose([
Resize((args.height,args.width)),
T.RandomHorizontalFlip(),
T.ToTensor(),
normalizer,
T.RandomErasing(probability=0.5, sh=0.2, r1=0.3)
])
# Start training
iter_nums = args.iteration
start_epoch = args.start_epoch
cluster_list = []
top_percent = args.rho
EF = 100 // iter_nums + 1
eug = None
for iter_n in range(start_epoch, iter_nums):
#### get source datas' feature
if args.load_dist and iter_n == 0:
dist = pickle.load(open('dist' + str(args.num_split) + '.pkl', 'rb'))
euclidean_dist_list = dist['euclidean']
rerank_dist_list = dist['rerank']
else:
source_features, _ = extract_features(model, src_extfeat_loader, for_eval=False)
if isinstance(source_features[src_dataset.trainval[0][0]], list):
len_f = len(source_features[src_dataset.trainval[0][0]])
source_features = [torch.cat([source_features[f][i].unsqueeze(0) for f, _, _ in src_dataset.trainval], 0) for i in range(len_f)]
else:
source_features = torch.cat([source_features[f].unsqueeze(0) for f, _, _ in src_dataset.trainval], 0) # synchronization feature order with s_dataset.trainval
#### extract training images' features
print('Iteration {}: Extracting Target Dataset Features...'.format(iter_n+1))
target_features, _ = extract_features(model, tgt_extfeat_loader, for_eval=False)
if isinstance(target_features[tgt_dataset.trainval[0][0]], list):
len_f = len(target_features[tgt_dataset.trainval[0][0]])
target_features = [torch.cat([target_features[f][i].unsqueeze(0) for f, _, _ in tgt_dataset.trainval], 0) for i in range(len_f)]
else:
target_features = torch.cat([target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval], 0) # synchronization feature order with dataset.trainval
#### calculate distance and rerank result
print('Calculating feature distances...')
# target_features = target_features.numpy()
euclidean_dist_list, rerank_dist_list = compute_dist(
source_features, target_features, lambda_value=args.lambda_value, no_rerank=args.no_rerank, num_split=args.num_split) # lambda=1 means only source dist
del target_features
del source_features
labels_list, cluster_list = generate_selflabel(
euclidean_dist_list, rerank_dist_list, iter_n, args, cluster_list)
#### generate new dataset
train_loader = generate_dataloader(tgt_dataset, labels_list, train_transformer, iter_n, args)
if iter_n == 5:
u_data, l_data = updata_lable(tgt_dataset, labels_list[0], args.tgt_dataset, sample=args.sample)
eug = EUG(model_name=args.arch, batch_size=args.batch_size, mode=args.mode, num_classes=num_class,
data_dir=args.data_dir, l_data=l_data, u_data=u_data, print_freq=args.print_freq,
save_path=args.logs_dir, pretrained_model=model, rerank=True)
eug.model = model
if eug is not None:
nums_to_select = int(min((iter_n + 1) * int(len(u_data) // (iter_nums)), len(u_data)))
pred_y, pred_score = eug.estimate_label()
print('This is running {} with EF= {}%, step {}:\t Nums_to_be_select {}, \t Logs-dir {}'.format(
args.mode, EF, iter_n+1, nums_to_select, args.logs_dir
))
selected_idx = eug.select_top_data(pred_score, nums_to_select)
new_train_data = eug.generate_new_train_data(selected_idx, pred_y)
eug_dataloader = eug.get_dataloader(new_train_data, training=True)
top1 = iter_trainer(model, tgt_dataset, train_loader, eug_dataloader, test_loader, optimizer,
criterion, args.epochs, args.logs_dir, args.print_freq, args.lr)
eug.model = model
del train_loader
# del eug_dataloader
else:
top1 = iter_trainer(model, tgt_dataset, train_loader, None, test_loader, optimizer,
criterion, args.epochs, args.logs_dir, args.print_freq, args.lr)
del train_loader
is_best = top1 > best_top1
best_top1 = max(top1, best_top1)
save_checkpoint({
'state_dict': model.module.state_dict(),
'epoch': iter_n + 1,
'best_top1': best_top1,
# 'num_ids': num_ids,
}, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'.
format(iter_n+1, top1, best_top1, ' *' if is_best else ''))
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch == 'inception' else \
(384, 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, args.num_instances, args.workers,
args.combine_trainval)
# Create model
# Hacking here to let the classifier be the last feature embedding layer
# Net structure: avgpool -> FC(1024) -> FC(args.features)
base_model = models.create(args.arch,
num_features=1024,
cut_at_pooling=True,
dropout=args.dropout,
num_classes=args.features)
grp_num = args.grp_num
embed_model = [
VNetEmbed(instances_num=args.num_instances,
feat_num=(2048 / grp_num),
num_classes=2,
drop_ratio=args.dropout).cuda() for i in range(grp_num)
]
base_model = nn.DataParallel(base_model).cuda()
model = VNetExtension(
instances_num=args.num_instances, #
base_model=base_model,
embed_model=embed_model,
alpha=args.alpha)
if args.retrain:
if args.evaluate_from:
print('loading trained model...')
checkpoint = load_checkpoint(args.evaluate_from)
model.load_state_dict(checkpoint['state_dict'])
else:
print('loading base part of pretrained model...')
checkpoint = load_checkpoint(args.retrain)
#copy_state_dict(checkpoint['state_dict'], base_model, strip='base.module.', replace='module.')
copy_state_dict(checkpoint['state_dict'],
base_model,
strip='base_model.',
replace='')
print('loading embed part of pretrained model...')
if grp_num > 1:
for i in range(grp_num):
copy_state_dict(checkpoint['state_dict'],
embed_model[i],
strip='embed_model.bn_' + str(i) + '.',
replace='bn.')
copy_state_dict(checkpoint['state_dict'],
embed_model[i],
strip='embed_model.classifier_' + str(i) +
'.',
replace='classifier.')
else:
copy_state_dict(checkpoint['state_dict'],
embed_model[0],
strip='module.embed_model.',
replace='')
# Distance metric
metric = DistanceMetric(algorithm=args.dist_metric)
# Load from checkpoint
start_epoch = best_top1 = 0
best_mAP = 0
if args.resume:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
best_top1 = checkpoint['best_top1']
print("=> Start epoch {} best top1 {:.1%}".format(
start_epoch, best_top1))
# Evaluator
evaluator = CascadeEvaluator(
base_model,
embed_model,
embed_dist_fn=lambda x: F.softmax(x, dim=1).data[:, 0])
#embed_dist_fn=lambda x: F.softmax(x))# here we are performing softmax normalization, this function take N,2 vector and after normalizing both column it return the
#first column
if args.evaluate:
metric.train(model, train_loader)
if args.evaluate_from:
print('loading trained model...')
checkpoint = load_checkpoint(args.evaluate_from)
model.load_state_dict(checkpoint['state_dict'])
print("Test:")
evaluator.evaluate(test_loader,
dataset.query,
dataset.gallery,
args.alpha,
metric,
rerank_topk=args.rerank,
dataset=args.dataset)
return
# Criterion
criterion = nn.CrossEntropyLoss().cuda()
criterion2 = TripletLoss(margin=args.margin).cuda()
#criterion = nn.BCELoss().cuda()
# base lr rate and embed lr rate
new_params = [z for z in model.embed]
param_groups = [
{'params': model.base.module.base.parameters(), 'lr_mult': 1.0}] + \
[{'params': new_params[i].parameters(), 'lr_mult': 10.0} for i in range(grp_num)]
# Optimizer
optimizer = torch.optim.Adam(param_groups,
lr=args.lr,
weight_decay=args.weight_decay)
# Trainer
trainer = DCDSBase(model, criterion, criterion2, args.alpha, grp_num)
# Schedule learning rate
def adjust_lr(epoch):
step_size = args.ss if args.arch == 'inception' else 20
lr = args.lr * (0.1**(epoch // step_size))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
return lr
# Start training
for epoch in range(start_epoch, args.epochs):
lr = adjust_lr(epoch)
trainer.train(epoch, train_loader, optimizer, lr, warm_up=False)
top1, mAP = evaluator.evaluate(val_loader,
dataset.val,
dataset.val,
args.alpha,
rerank_topk=args.rerank,
second_stage=True,
dataset=args.dataset)
is_best = top1 > best_top1
best_mAP = max(mAP, best_mAP)
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': epoch + 1,
'best_top1': best_top1,
},
is_best,
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print('\n * Finished epoch {:3d} mAP: {:5.1%} best: {:5.1%}{}\n'.
format(epoch, mAP, best_mAP, ' *' 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'])
metric.train(model, train_loader)
evaluator.evaluate(test_loader,
dataset.query,
dataset.gallery,
args.alpha,
metric,
rerank_topk=args.rerank,
dataset=args.dataset)
def main(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
if args.height is None or args.width is None:
args.height, args.width = (144, 56) if args.arch == 'inception' else \
(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, args.num_instances, args.workers,
args.combine_trainval, args.batch_id)
# Create model
# Hacking here to let the classifier be the last feature embedding layer
# Net structure: avgpool -> FC(1024) -> FC(args.features)
model = models.create(args.arch,
num_features=1024,
dropout=args.dropout,
num_classes=args.features)
# Load from checkpoint
start_epoch = best_top1 = 0
if args.resume:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
best_top1 = checkpoint['best_top1']
print("=> Start epoch {} best top1 {:.1%}".format(
start_epoch, best_top1))
model = nn.DataParallel(model).cuda()
# Distance metric
metric = DistanceMetric(algorithm=args.dist_metric)
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
metric.train(model, train_loader)
print("Validation:")
evaluator.evaluate(val_loader, dataset.val, dataset.val, metric)
print("Test:")
evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
return
# Criterion
criterion = TripletLoss(margin=args.margin).cuda()
# Optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# Trainer
trainer = Trainer(model, criterion)
# Schedule learning rate
def adjust_lr(epoch):
lr = args.lr if epoch <= 100 else \
args.lr * (0.001 ** ((epoch - 100) / 50.0))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
# Start training
for epoch in range(start_epoch, args.epochs):
adjust_lr(epoch)
trainer.train(epoch, train_loader, optimizer)
if epoch < args.start_save:
continue
top1 = evaluator.evaluate(val_loader, dataset.val, dataset.val)
is_best = top1 > best_top1
best_top1 = max(top1, best_top1)
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch + 1,
'best_top1': best_top1,
},
is_best,
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'.
format(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.module.load_state_dict(checkpoint['state_dict'])
metric.train(model, train_loader)
evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
def main(args):
print(args)
checkpoint_save_dir = os.path.join(args.logs_dir,
'process_%d' % (args.i % args.nums))
if not os.path.exists(checkpoint_save_dir):
os.makedirs(checkpoint_save_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
if not args.evaluate:
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Create data loaders
if args.big_height is None or args.big_width is None or args.target_height is None or args.target_width is None:
args.big_height, args.big_width, args.target_height, args.target_width = (
256, 256, 224, 224)
dataset, num_classes, train_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.big_height, args.big_width,
args.target_height, args.target_width, args.batch_size, args.num_instances,
args.workers, args.combine_trainval)
# Create model
model = models.create(args.arch,
num_classes=num_classes,
num_features=args.features)
print(model.state_dict().keys())
# Load from checkpoint
start_epoch = best = 0
if args.weights and args.arch != 'densenet161':
checkpoint = load_checkpoint(args.weights)
if args.arch == 'cross_trihard_senet101' or args.arch == 'cross_trihard_se_resnet152':
del (checkpoint['last_linear.weight'])
del (checkpoint['last_linear.bias'])
model.base.load_state_dict(checkpoint)
#model.base.load_param(args.weights)
elif args.arch == 'cross_trihard_mobilenet':
del (checkpoint['state_dict']['module.fc.weight'])
del (checkpoint['state_dict']['module.fc.bias'])
model_dict = model.state_dict()
checkpoint_load = {
k.replace('module', 'base'): v
for k, v in (checkpoint['state_dict']).items()
}
model_dict.update(checkpoint_load)
model.load_state_dict(model_dict)
elif args.arch == 'cross_trihard_shufflenetv2':
del (checkpoint['classifier.0.weight'])
del (checkpoint['classifier.0.bias'])
model.base.load_state_dict(checkpoint)
elif args.arch == 'cross_trihard_densenet121':
del (checkpoint['classifier.weight'])
del (checkpoint['classifier.bias'])
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$'
)
for key in list(checkpoint.keys()):
res = pattern.match(key)
if res:
new_key = res.group(1) + res.group(2)
checkpoint[new_key] = checkpoint[key]
del checkpoint[key]
model.base.load_state_dict(checkpoint)
elif args.arch == 'cross_trihard_vgg19bn':
del (checkpoint['classifier.6.weight'])
del (checkpoint['classifier.6.bias'])
model.base.load_state_dict(checkpoint)
else:
del (checkpoint['fc.weight'])
del (checkpoint['fc.bias'])
model.base.load_state_dict(checkpoint)
if args.resume and not args.weights:
checkpoint = load_checkpoint(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
print("=> Start epoch {} best top1 {:.1%}".format(start_epoch, best))
model = nn.DataParallel(model).cuda()
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
print("Test:")
evaluator.evaluate(test_loader, dataset.query, dataset.gallery)
return
# Criterion
ranking_loss = nn.MarginRankingLoss(margin=args.margin).cuda()
criterion = { 'crossentropy': nn.CrossEntropyLoss().cuda(), \
'trihard': TripletLoss(ranking_loss).cuda() }
# criterion = { 'crossentropy': nn.CrossEntropyLoss().cuda(), \
# 'trihard': TripletLoss(ranking_loss).cuda(), \
# 'center': CenterLoss(num_classes= num_classes,feat_dim= args.features) }
# criterion = { 'crossentropy':nn.crossentropy.cuda(), \
# 'trihard': TripletLoss(ranking_loss).cuda() }
# Optimizer
if hasattr(model.module, 'base'):
base_params = []
base_bn_params = []
for name, p in model.module.base.named_parameters():
if 'bn' in name:
base_bn_params.append(p)
else:
base_params.append(p)
base_param_ids = set(map(id, model.module.base.parameters()))
new_params = [
p for p in model.parameters() if id(p) not in base_param_ids
]
param_groups = [{
'params': base_params,
'lr_mult': 0.1
}, {
'params': base_bn_params,
'lr_mult': 0.1
}, {
'params': new_params,
'lr_mult': args.lr_mult
}]
else:
param_groups = model.parameters()
if args.optimizer == 0:
optimizer = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
else:
print('Adam')
optimizer = torch.optim.Adam(params=param_groups,
lr=args.lr,
weight_decay=args.weight_decay)
# Trainer
trainer = Cross_Trihard_Trainer(model,
criterion,
metric_loss_weight=args.metric_loss_weight)
# Schedule learning rate
def adjust_lr(epoch):
step_size, step_size2, step_size3 = args.step_size, args.step_size2, args.step_size3
#lr = args.lr * (0.1 ** (epoch // step_size))
if epoch <= step_size:
lr = args.lr
elif epoch <= step_size2:
lr = args.lr * 0.1
elif epoch <= step_size3:
lr = args.lr * 0.01
else:
lr = args.lr * 0.001
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
return lr
# Start training
for epoch in range(start_epoch + 1, args.epochs + 1):
lr = adjust_lr(epoch)
trainer.train(epoch,
train_loader,
optimizer,
lr,
warm_up=True,
warm_up_ep=args.warm_up_ep,
accumu_step=1)
if epoch % args.epoch_inter == 0 or epoch >= args.dense_evaluate:
tmp_res = evaluator.evaluate(test_loader, dataset.query,
dataset.gallery)
print('tmp_res: ', tmp_res)
print('best: ', best)
if epoch >= args.start_save:
save_checkpoint(
{
'state_dict': model.module.state_dict(),
'epoch': epoch,
},
False,
fpath=osp.join(
args.logs_dir, 'process_%d/epoch_%d.pth.tar' %
(args.i % args.nums, epoch)))
if tmp_res > best:
shutil.copy(
osp.join(
args.logs_dir, 'process_%d/epoch_%d.pth.tar' %
(args.i % args.nums, epoch)),
osp.join(
args.logs_dir, 'process_%d_model_best.pth.tar' %
(args.i % args.nums)))
best = tmp_res
outside(model, epoch, args)