def create_model(args):
model_1 = models.create(args.arch,
num_features=args.features,
dropout=args.dropout,
num_classes=1)
model_1_ema = models.create(args.arch,
num_features=args.features,
dropout=args.dropout,
num_classes=1)
model_1.cuda()
model_1_ema.cuda()
model_1 = nn.DataParallel(model_1)
model_1_ema = nn.DataParallel(model_1_ema)
if args.no_source:
print('No source pre-training')
else:
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)
for param in model_1_ema.parameters():
param.detach_()
return model_1, model_1_ema
def create_model(args, classes, phase):
'''
创建模型，新建模型与保存的模型结构相同，以便可以进行载入权重
:param args:
:param classes:
:param phase:
:return:
'''
model_ema = models.create(args.arch,
num_features=args.features,
dropout=args.dropout,
num_classes=classes)
initial_weights = load_checkpoint(
osp.join(args.init, 'phase{}_model_best.pth.tar'.format(phase)))
if 'module.classifier.fc1.weight' in initial_weights['state_dict'].keys():
in_features = initial_weights['state_dict'][
'module.classifier.fc1.weight'].data.size(1)
out_features1 = initial_weights['state_dict'][
'module.classifier.fc1.weight'].data.size(0)
out_features2 = initial_weights['state_dict'][
'module.classifier.fc2.weight'].data.size(0)
new_fc = SplitCosineLinear(in_features,
out_features1,
out_features2,
sigma=True)
model_ema.classifier = new_fc
new_fc_max = copy.deepcopy(new_fc)
model_ema.classifier_max = new_fc_max
else:
in_features = initial_weights['state_dict'][
'module.classifier.weight'].data.size(1)
out_features = initial_weights['state_dict'][
'module.classifier.weight'].data.size(0)
new_fc = CosineLinear(in_features=in_features,
out_features=out_features,
sigma=True)
model_ema.classifier = new_fc
new_fc_max = copy.deepcopy(new_fc)
model_ema.classifier_max = new_fc_max
copy_state_dict(initial_weights['state_dict'], model_ema, strip='module.')
model_cur = copy.deepcopy(model_ema)
model_ref = copy.deepcopy(model_ema)
model_ema.cuda()
model_ema = nn.DataParallel(model_ema)
model_cur.cuda()
model_cur = nn.DataParallel(model_cur)
model_ref.cuda()
model_ref = nn.DataParallel(model_ref)
for param in model_ema.parameters():
param.detach_()
return model_cur, model_ema, model_ref
def main_worker(args):
global start_epoch, 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)
ori_train = dataset_target.train
if not args.no_source:
dataset_source = get_data(args.dataset_source, args.data_dir)
test_loader_target = get_test_loader(dataset_target, args.height, args.width, args.batch_size, args.workers)
# Create model
model_1, model_1_ema = create_model(args)
# Evaluator
evaluator_1_ema = Evaluator(model_1_ema)
best_mAP = 0
for nc in range(args.epochs):
cluster_loader = get_test_loader(dataset_target, args.height, args.width, args.batch_size, args.workers,
testset=dataset_target.train)
dict_f, _ = extract_features(model_1_ema, cluster_loader, print_freq=50)
cf_1 = torch.stack(list(dict_f.values()))
# DBSCAN cluster
if args.no_source:
rerank_dist = compute_jaccard_dist(cf_1, lambda_value=0, source_features=None,
use_gpu=False).numpy()
else:
cluster_loader_source = get_test_loader(dataset_source, args.height, args.width, args.batch_size,
args.workers, testset=dataset_source.train)
dict_f_source, _ = extract_features(model_1_ema, cluster_loader_source, print_freq=50)
cf_1_source = torch.stack(list(dict_f_source.values()))
rerank_dist = compute_jaccard_dist(cf_1, lambda_value=args.lambda_value, source_features=cf_1_source,
use_gpu=False).numpy()
del cf_1_source
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))
print('Clustering and labeling...')
cluster = DBSCAN(eps=eps, min_samples=4, metric='precomputed', n_jobs=-1)
labels = cluster.fit_predict(rerank_dist)
num_ids = len(set(labels)) -1
print('Epoch {} have {} training ids'.format(nc, num_ids))
# generate new dataset
labeled_ind, unlabeled_ind = [], []
for ind, label in enumerate(labels):
if label == -1:
unlabeled_ind.append(ind)
else:
labeled_ind.append(ind)
# print('Epoch {} have {} labeled samples and {} unlabeled samples'.format(nc + 1, len(labeled_ind), len(unlabeled_ind)))
cf_1 = cf_1.numpy()
centers = []
for id in range(num_ids):
centers.append(np.mean(cf_1[labels == id], axis=0))
centers = np.stack(centers, axis=0)
del cf_1, rerank_dist
model_1.module.classifier = nn.Linear(2048, num_ids, bias=False).cuda()
model_1_ema.module.classifier = nn.Linear(2048, num_ids, bias=False).cuda()
model_1.module.classifier_max = nn.Linear(2048, num_ids, bias=False).cuda()
model_1_ema.module.classifier_max = nn.Linear(2048, num_ids, bias=False).cuda()
model_1.module.classifier.weight.data.copy_(
torch.from_numpy(normalize(centers[:, :2048], axis=1)).float().cuda())
model_1_ema.module.classifier.weight.data.copy_(
torch.from_numpy(normalize(centers[:, :2048], axis=1)).float().cuda())
model_1.module.classifier_max.weight.data.copy_(
torch.from_numpy(normalize(centers[:, 2048:], axis=1)).float().cuda())
model_1_ema.module.classifier_max.weight.data.copy_(
torch.from_numpy(normalize(centers[:, 2048:], axis=1)).float().cuda())
del centers
target_label = 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])
# Optimizer
params = []
for key, value in model_1.named_parameters():
if not value.requires_grad:
continue
params += [{"params": [value], "lr": args.lr, "weight_decay": args.weight_decay}]
optimizer = torch.optim.Adam(params)
# Trainer
trainer = ABMTTrainer(model_1, model_1_ema, num_cluster=num_ids, alpha=args.alpha)
epoch = nc
# # DBSCAN
dataset_target.train = [ori_train[i] for i in labeled_ind]
print(len(dataset_target.train), 'are labeled.')
labeled_loader_target = get_train_loader(dataset_target, args.height, args.width,
args.batch_size, args.workers, args.num_instances, iters, mutual=True)
labeled_loader_target.new_epoch()
trainer.train(epoch, labeled_loader_target, optimizer,
print_freq=args.print_freq, train_iters=len(labeled_loader_target))
def save_model(model_ema, is_best, best_mAP, mid, num_ids):
save_checkpoint({
'state_dict': model_ema.state_dict(),
'epoch': epoch + 1,
'best_mAP': best_mAP,
'num_ids': num_ids
}, 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)):
print('Evaluating teacher net:')
cmc, mAP_1 = evaluator_1_ema.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery, cmc_flag=True)
is_best = (mAP_1>best_mAP)
best_mAP = max(mAP_1, best_mAP)
save_model(model_1_ema, is_best, best_mAP, 1, num_ids)
dataset_target.train = ori_train
print ('Test on the best model.')
checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
model_best = models.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=checkpoint['num_ids'])
model_best.cuda()
model_best = nn.DataParallel(model_best)
evaluator_best = Evaluator(model_best)
model_best.load_state_dict(checkpoint['state_dict'])
evaluator_best.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.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 = ABMTPreTrainer(model, num_classes, margin=args.margin)
# Start training
for epoch in range(start_epoch, args.epochs):
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)
lr_scheduler.step()
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)