def create_model(args):
model_1 = models_pre_a_s2.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=args.num_clusters)
model_2 = models_pre_a_s2.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=args.num_clusters)
model_1_ema = models_pre_a_s2.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=args.num_clusters)
model_2_ema = models_pre_a_s2.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=args.num_clusters)
model_1.cuda()
model_2.cuda()
model_1_ema.cuda()
model_2_ema.cuda()
model_1 = nn.DataParallel(model_1)
model_2 = nn.DataParallel(model_2)
model_1_ema = nn.DataParallel(model_1_ema)
model_2_ema = nn.DataParallel(model_2_ema)
initial_weights = load_checkpoint(args.init_1)
copy_state_dict(initial_weights['state_dict'], model_1)
copy_state_dict(initial_weights['state_dict'], model_1_ema)
model_1_ema.module.classifier.weight.data.copy_(model_1.module.classifier.weight.data)
initial_weights = load_checkpoint(args.init_2)
copy_state_dict(initial_weights['state_dict'], model_2)
copy_state_dict(initial_weights['state_dict'], model_2_ema)
model_2_ema.module.classifier.weight.data.copy_(model_2.module.classifier.weight.data)
for param in model_1_ema.parameters():
param.detach_()
for param in model_2_ema.parameters():
param.detach_()
return model_1, model_2, model_1_ema, model_2_ema
def main_worker(args):
cudnn.benchmark = True
log_dir = osp.dirname(args.resume)
sys.stdout = Logger(osp.join(log_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
# Create data loaders
dataset_target, test_loader_target = \
get_data(args.dataset_target, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
model = models_post_a_s1.create(args.arch,
pretrained=False,
num_features=args.features,
dropout=args.dropout,
num_classes=0)
model.cuda()
model = nn.DataParallel(model)
# Load from checkpoint
checkpoint = load_checkpoint(args.resume)
copy_state_dict(checkpoint['state_dict'], model)
start_epoch = checkpoint['epoch']
best_mAP = checkpoint['best_mAP']
print("=> Checkpoint of epoch {} best mAP {:.1%}".format(
start_epoch, best_mAP))
# Evaluator
evaluator = Evaluator(model)
print("Test on the target domain of {}:".format(args.dataset_target))
evaluator.evaluate(test_loader_target,
dataset_target.query,
dataset_target.gallery,
cmc_flag=True,
rerank=args.rerank)
return
def main_worker(args):
global best_mAP
cudnn.benchmark = True
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
print("==========\nArgs:{}\n==========".format(args))
# Create data loaders
iters = args.iters if (args.iters > 0) else None
dataset_target = get_data(args.dataset_target, args.data_dir)
test_loader_target = get_test_loader(dataset_target, args.height,
args.width, args.batch_size,
args.workers)
cluster_loader = get_test_loader(dataset_target,
args.height,
args.width,
args.batch_size,
args.workers,
testset=dataset_target.train)
# Create model
model_1, model_2, model_1_ema, model_2_ema = create_model(args)
# Evaluator
evaluator_1_ema = Evaluator(model_1_ema)
evaluator_2_ema = Evaluator(model_2_ema)
for epoch in range(args.epochs):
dict_f, _ = extract_features(model_1_ema,
cluster_loader,
print_freq=50)
cf_1 = torch.stack(list(dict_f.values())).numpy()
dict_f, _ = extract_features(model_2_ema,
cluster_loader,
print_freq=50)
cf_2 = torch.stack(list(dict_f.values())).numpy()
cf = (cf_1 + cf_2) / 2
print('\n Clustering into {} classes \n'.format(args.num_clusters))
km = KMeans(n_clusters=args.num_clusters,
random_state=args.seed,
n_jobs=2).fit(cf)
model_1.module.classifier.weight.data.copy_(
torch.from_numpy(normalize(km.cluster_centers_,
axis=1)).float().cuda())
model_2.module.classifier.weight.data.copy_(
torch.from_numpy(normalize(km.cluster_centers_,
axis=1)).float().cuda())
model_1_ema.module.classifier.weight.data.copy_(
torch.from_numpy(normalize(km.cluster_centers_,
axis=1)).float().cuda())
model_2_ema.module.classifier.weight.data.copy_(
torch.from_numpy(normalize(km.cluster_centers_,
axis=1)).float().cuda())
target_label = km.labels_
# change pseudo labels
for i in range(len(dataset_target.train)):
dataset_target.train[i] = list(dataset_target.train[i])
dataset_target.train[i][1] = int(target_label[i])
dataset_target.train[i] = tuple(dataset_target.train[i])
train_loader_target = get_train_loader(dataset_target, args.height,
args.width, args.batch_size,
args.workers,
args.num_instances, iters)
# Optimizer
params = []
for key, value in model_1.named_parameters():
if not value.requires_grad:
continue
if 'non_local' in key:
params += [{
"params": [value],
"lr": args.lr,
"weight_decay": args.weight_decay
}]
for key, value in model_2.named_parameters():
if not value.requires_grad:
continue
if 'non_local' in key:
params += [{
"params": [value],
"lr": args.lr,
"weight_decay": args.weight_decay
}]
optimizer = torch.optim.Adam(params)
# Trainer
trainer = MMTTrainer(model_1,
model_2,
model_1_ema,
model_2_ema,
num_cluster=args.num_clusters,
alpha=args.alpha)
train_loader_target.new_epoch()
trainer.train(epoch,
train_loader_target,
optimizer,
ce_soft_weight=args.soft_ce_weight,
tri_soft_weight=args.soft_tri_weight,
print_freq=args.print_freq,
train_iters=len(train_loader_target))
def save_model(model_ema, is_best, best_mAP, mid):
save_checkpoint(
{
'state_dict': model_ema.state_dict(),
'epoch': epoch + 1,
'best_mAP': best_mAP,
},
is_best,
fpath=osp.join(args.logs_dir,
'model' + str(mid) + '_checkpoint.pth.tar'))
if ((epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1)):
mAP_1 = evaluator_1_ema.evaluate(test_loader_target,
dataset_target.query,
dataset_target.gallery,
cmc_flag=False)
mAP_2 = evaluator_2_ema.evaluate(test_loader_target,
dataset_target.query,
dataset_target.gallery,
cmc_flag=False)
is_best = (mAP_1 > best_mAP) or (mAP_2 > best_mAP)
best_mAP = max(mAP_1, mAP_2, best_mAP)
save_model(model_1_ema, (is_best and (mAP_1 > mAP_2)), best_mAP, 1)
save_model(model_2_ema, (is_best and (mAP_1 <= mAP_2)), best_mAP,
2)
print(
'\n * Finished epoch {:3d} model no.1 mAP: {:5.1%} model no.2 mAP: {:5.1%} best: {:5.1%}{}\n'
.format(epoch, mAP_1, mAP_2, best_mAP,
' *' if is_best else ''))
print('Test on the best model.')
checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
model_1_ema.load_state_dict(checkpoint['state_dict'])
evaluator_1_ema.evaluate(test_loader_target,
dataset_target.query,
dataset_target.gallery,
cmc_flag=True)
def main_worker(args):
global start_epoch, best_mAP
cudnn.benchmark = True
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
iters = args.iters if (args.iters > 0) else None
dataset_source, num_classes, train_loader_source, test_loader_source = \
get_data(args.dataset_source, args.data_dir, args.height,
args.width, args.batch_size, args.workers, args.num_instances, iters)
dataset_target, _, train_loader_target, test_loader_target = \
get_data(args.dataset_target, args.data_dir, args.height,
args.width, args.batch_size, args.workers, 0, iters)
# Create model
model = models_source.create(args.arch,
num_features=args.features,
dropout=args.dropout,
num_classes=num_classes)
model.cuda()
model = nn.DataParallel(model)
# Load from checkpoint
if args.resume:
checkpoint = load_checkpoint(args.resume)
copy_state_dict(checkpoint['state_dict'], model)
start_epoch = checkpoint['epoch']
best_mAP = checkpoint['best_mAP']
print("=> Start epoch {} best mAP {:.1%}".format(
start_epoch, best_mAP))
# Evaluator
evaluator = Evaluator(model)
if args.evaluate:
print("Test on source domain:")
evaluator.evaluate(test_loader_source,
dataset_source.query,
dataset_source.gallery,
cmc_flag=True,
rerank=args.rerank)
print("Test on target domain:")
evaluator.evaluate(test_loader_target,
dataset_target.query,
dataset_target.gallery,
cmc_flag=True,
rerank=args.rerank)
return
params = []
for key, value in model.named_parameters():
if not value.requires_grad:
continue
params += [{
"params": [value],
"lr": args.lr,
"weight_decay": args.weight_decay
}]
optimizer = torch.optim.Adam(params)
lr_scheduler = WarmupMultiStepLR(optimizer,
args.milestones,
gamma=0.1,
warmup_factor=0.01,
warmup_iters=args.warmup_step)
# Trainer
trainer = PreTrainer(model, num_classes, margin=args.margin)
# Start training
for epoch in range(start_epoch, args.epochs):
lr_scheduler.step()
train_loader_source.new_epoch()
train_loader_target.new_epoch()
trainer.train(epoch,
train_loader_source,
train_loader_target,
optimizer,
train_iters=len(train_loader_source),
print_freq=args.print_freq)
if ((epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1)):
_, mAP = evaluator.evaluate(test_loader_source,
dataset_source.query,
dataset_source.gallery,
cmc_flag=True)
is_best = mAP > best_mAP
best_mAP = max(mAP, best_mAP)
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': epoch + 1,
'best_mAP': best_mAP,
},
is_best,
fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print(
'\n * Finished epoch {:3d} source mAP: {:5.1%} best: {:5.1%}{}\n'
.format(epoch, mAP, best_mAP, ' *' if is_best else ''))
print("Test on target domain:")
evaluator.evaluate(test_loader_target,
dataset_target.query,
dataset_target.gallery,
cmc_flag=True,
rerank=args.rerank)
