def main_worker(confg):
global start_epoch, best_mAP
cudnn.benchmark = True #
sys.stdout = Logger(os.path.join(config['logs_dir'], 'log.txt'))
print("==========\nArgs:{}\n==========".format(config))
#######################################################
# create test data loaders
#######################################################
if config['iters'] > 0:
iters = config['iters']
else:
iters = None
#dataset_target = get_data(config['dataset_target'], config['data_dir'])
test_loader_target = get_test_loader(config['data_dir'],
config['resize_size'],
config['batch_size'],
config['workers'])
#######################################################
# create model
#######################################################
model_1, model_2, model_1_ema, model_2_ema = create_model(config)
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
query_loader, gallery_loader = \
get_data(args.dataset_target, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
model = models.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
q_f, q_id, q_cam = test(model, query_loader)
g_f, g_id, g_cam = test(model, gallery_loader)
q_g_dist = np.dot(q_f, np.transpose(g_f))
q_g_dist = 2. - 2 * q_g_dist # change the cosine similarity metric to euclidean similarity metric
all_cmc, mAP = eval_func(q_g_dist, q_id, g_id, q_cam, g_cam)
all_cmc = all_cmc * 100
print('rank-1: {:.4f} rank-5: {:.4f} rank-10: {:.4f} rank-20: {:.4f} rank-50: {:.4f} mAP: {:.4f}'.format(all_cmc[0],
all_cmc[4],
all_cmc[9],
all_cmc[
19],
all_cmc[
49],
mAP * 100))
indices = np.argsort(q_g_dist, axis=1)
np.savetxt("answer.txt", indices[:, :100], fmt="%04d")
return
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.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 = create_model(args)
# Evaluator
evaluator = Evaluator(model)
for epoch in range(args.epochs):
dict_f, _ = extract_features(model, cluster_loader, print_freq=50)
cf = torch.stack(list(dict_f.values())).numpy()
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.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.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 = ClusterBaseTrainer(model, num_cluster=args.num_clusters)
train_loader_target.new_epoch()
trainer.train(epoch,
train_loader_target,
optimizer,
print_freq=args.print_freq,
train_iters=len(train_loader_target))
def save_model(model, is_best, 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'))
if (epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1):
mAP = evaluator.evaluate(test_loader_target,
dataset_target.query,
dataset_target.gallery,
cmc_flag=False)
is_best = (mAP > best_mAP)
best_mAP = max(mAP, best_mAP)
save_model(model, is_best, best_mAP)
print(
'\n * Finished epoch {:3d} model mAP: {:5.1%} best: {:5.1%}{}\n'
.format(epoch, mAP, 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.load_state_dict(checkpoint['state_dict'])
evaluator.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 = \
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 = \
get_data(args.dataset_target, args.data_dir, args.height,
args.width, args.batch_size, args.workers, args.num_instances, iters)
dataset_validation, test_loader_target = \
get_test_data(args.dataset_validation, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
model = models.create(args.arch,
dropout=args.dropout,
num_classes=num_classes,
circle=args.circle)
# print(model)
print("Model size: {:.3f} M".format(count_num_param(model)))
model.cuda()
model = nn.DataParallel(model) # 多gpu并行
# 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, args, 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,
optimizer,
train_iters=len(train_loader_source),
print_freq=args.print_freq,
balance=args.balance)
if (epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1):
_, mAP = evaluator.evaluate(test_loader_target,
dataset_validation.query,
dataset_validation.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_validation.query,
dataset_validation.gallery,
cmc_flag=True,
rerank=args.rerank)
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_source = get_data(args.dataset_source, args.data_dir)
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)
tar_cluster_loader = get_test_loader(dataset_target,
args.height,
args.width,
args.batch_size,
args.workers,
testset=dataset_target.train)
sour_cluster_loader = get_test_loader(dataset_source,
args.height,
args.width,
args.batch_size,
args.workers,
testset=dataset_source.train)
# Create model
model_1, model_2, model_1_ema, model_2_ema = create_model(
args, len(dataset_target.train))
# 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,
tar_cluster_loader,
print_freq=50)
cf_1 = torch.stack(list(dict_f.values()))
dict_f, _ = extract_features(model_2_ema,
tar_cluster_loader,
print_freq=50)
cf_2 = torch.stack(list(dict_f.values()))
cf = (cf_1 + cf_2) / 2
cf = F.normalize(cf, dim=1)
if (args.lambda_value > 0):
dict_f, _ = extract_features(model_1_ema,
sour_cluster_loader,
print_freq=50)
cf_1 = torch.stack(list(dict_f.values()))
dict_f, _ = extract_features(model_2_ema,
sour_cluster_loader,
print_freq=50)
cf_2 = torch.stack(list(dict_f.values()))
cf_s = (cf_1 + cf_2) / 2
cf_s = F.normalize(cf_s, dim=1)
rerank_dist = compute_jaccard_dist(cf,
lambda_value=args.lambda_value,
source_features=cf_s,
use_gpu=args.rr_gpu).numpy()
else:
rerank_dist = compute_jaccard_dist(cf, use_gpu=args.rr_gpu).numpy()
if (epoch == 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)
rho = 1.6e-3
top_num = np.round(rho * tri_mat.size).astype(int)
eps = tri_mat[:top_num].mean()
print('eps for cluster: {:.3f}'.format(eps))
cluster = DBSCAN(eps=eps,
min_samples=4,
metric='precomputed',
n_jobs=-1)
print('Clustering and labeling...')
labels = cluster.fit_predict(rerank_dist)
num_ids = len(set(labels)) - (1 if -1 in labels else 0)
args.num_clusters = num_ids
print('\n Clustered into {} classes \n'.format(args.num_clusters))
# generate new dataset and calculate cluster centers
new_dataset = []
cluster_centers = collections.defaultdict(list)
for i, ((fname, _, cid),
label) in enumerate(zip(dataset_target.train, labels)):
if label == -1: continue
new_dataset.append((fname, label, cid))
cluster_centers[label].append(cf[i])
cluster_centers = [
torch.stack(cluster_centers[idx]).mean(0)
for idx in sorted(cluster_centers.keys())
]
cluster_centers = torch.stack(cluster_centers)
model_1.module.classifier.weight.data[:args.num_clusters].copy_(
F.normalize(cluster_centers, dim=1).float().cuda())
model_2.module.classifier.weight.data[:args.num_clusters].copy_(
F.normalize(cluster_centers, dim=1).float().cuda())
model_1_ema.module.classifier.weight.data[:args.num_clusters].copy_(
F.normalize(cluster_centers, dim=1).float().cuda())
model_2_ema.module.classifier.weight.data[:args.num_clusters].copy_(
F.normalize(cluster_centers, dim=1).float().cuda())
train_loader_target = get_train_loader(dataset_target,
args.height,
args.width,
args.batch_size,
args.workers,
args.num_instances,
iters,
trainset=new_dataset)
# 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
}]
for key, value in model_2.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 = 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 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)
# test loader target 包含data set target的query和gallery
cluster_loader = get_test_loader(dataset_target, args.height, args.width, args.batch_size, args.workers,
testset=dataset_target.train)
# cluster loader 只包含data set 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) # 返回的是features labels的字典
cf_1 = torch.stack(list(dict_f.values())).numpy()
# dict_f.values()返回字典中的所有值
# list()将元组转换为列表
# torch.stack()沿着一个新维度对输入张量序列进行扩维连接，dim参数默认为0，形成n*x*y张量
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=-1).fit(cf)
# 聚类后的各个类的中心，将这些中心标准化以后直接赋值到了线性层的参数中
# 我们假设一个线性层的输入特征是M维的，输出是N维的，如果不加思考的化，就是一个M维到N维的转换。但是深入一点的思考的化，其实
# 可以将其理解为：将M维的向量，与N个M维的向量（Normalize后则是单位向量）做内积，生成了N个内积，内积i代表着输入特征在第i个
# 向量上的投影，这个值越大，代表着输入向量与这个向量的关系越紧密。
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
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
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, args)
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), balance=args.balance)
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
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)
# 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)
clusters = [args.num_clusters] * args.epochs
feature_length = args.features if args.features > 0 else 2048
moving_avg_features = np.zeros((len(dataset_target.train), feature_length))
for nc in range(len(clusters)):
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())).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
moving_avg_features = moving_avg_features * args.moving_avg_momentum + cf * (
1 - args.moving_avg_momentum)
moving_avg_features = moving_avg_features / (
1 - args.moving_avg_momentum**(nc + 1))
print('\n Clustering into {} classes \n'.format(clusters[nc]))
km = KMeans(n_clusters=clusters[nc], random_state=args.seed,
n_jobs=2).fit(moving_avg_features)
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
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
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=clusters[nc],
alpha=args.alpha)
train_loader_target.new_epoch()
epoch = nc
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)