def train_RRI(model, Ts: int=7):
base_lrs = [0.001] * 3 + [0.0001] * 10
for T in range(Ts):
print('=== T = {} ==='.format(T))
print('Replacing eigen layer weight...')
for eigen_layer in model.module.get_fcs():
replace_weight(eigen_layer)
print('Replaced.')
print('--- Restraint ({}) ---'.format(T))
train_R(model, base_lrs[T], T, fix_eigen_layer=True)
print('--- Relaxation ({}) ---'.format(T))
train_R(model, base_lrs[T], T, fix_eigen_layer=False)
for name in args.target_names:
print('Evaluating {} ...'.format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
rank1 = test(model, queryloader, galleryloader, use_gpu)
save_checkpoint({
'state_dict': model.state_dict(),
'rank1': rank1,
'epoch': 0,
'arch': args.arch,
'optimizer': (),
}, args.save_dir, prefix='final_')
def main():
global args
set_random_seed(1)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
log_name = 'test.log' if args.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(args.save_dir, log_name))
print('** Arguments **')
arg_keys = list(args.__dict__.keys())
arg_keys.sort()
for key in arg_keys:
print('{}: {}'.format(key, args.__dict__[key]))
torch.backends.cudnn.benchmark = True
datamanager = ImageDataManager(batch_size=args.batch_size)
trainloader, queryloader, galleryloader = datamanager.return_dataloaders()
print('Building model: {}'.format(args.arch))
model = build_model(args.arch,
4000,
args.bias,
args.bnneck,
pretrained=(not args.no_pretrained))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
model.cuda()
if args.evaluate:
evaluate(model, queryloader, galleryloader, args.dist_metric,
args.normalize_feature)
return
criterion = CrossEntropyLoss(4000)
optimizer = torch.optim.Adam(model.parameters(),
lr=0.0003,
weight_decay=5e-04,
betas=(0.9, 0.999))
scheduler = build_lr_scheduler(optimizer, args.lr_scheduler, args.stepsize)
time_start = time.time()
print('=> Start training')
for epoch in range(args.start_epoch, args.max_epoch):
train(epoch, model, criterion, optimizer, trainloader)
scheduler.step()
if (epoch + 1) % 20 == 0:
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': epoch + 1,
'optimizer': optimizer.state_dict(),
}, args.save_dir)
evaluate(model, queryloader, galleryloader, args.dist_metric,
args.normalize_feature)
elapsed = round(time.time() - time_start)
elapsed = str(datetime.timedelta(seconds=elapsed))
print('Elapsed {}'.format(elapsed))
def train_base(model):
use_sgd = os.environ.get('sgd') is not None
optimizer_getter = get_base_sgd_optimizer if use_sgd else get_base_optimizer
optimizer, scheduler = get_base_optimizer(model)
model.train()
print('=== train base ===')
if True:
open_layers = ['fc', 'classifier1', 'classifier2_1', 'classifier2_2', 'fc2_1', 'fc2_2', 'reduction', 'classifier']
print('Train {} for {} epochs while keeping other layers frozen'.format(open_layers, 10))
for epoch in range(10):
open_specified_layers(model, open_layers)
train(epoch, model, criterion, optimizer, trainloader, use_gpu, fixbase=True)
print('Done. All layers are open to train for {} epochs'.format(60))
open_all_layers(model)
optimizer, scheduler = optimizer_getter(model)
for epoch in range(60):
train(epoch, model, criterion, optimizer, trainloader, use_gpu=use_gpu)
scheduler.step()
print('=> Test')
if (epoch + 1) % args.eval_freq == 0:
for name in args.target_names:
print('Evaluating {} ...'.format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
rank1 = test(model, queryloader, galleryloader, use_gpu)
save_checkpoint({
'state_dict': model.state_dict(),
'rank1': rank1,
'epoch': 0,
'arch': args.arch,
'optimizer': optimizer.state_dict(),
}, args.save_dir, prefix='base_')
def train_R(model, lr, T, fix_eigen_layer: bool=False):
eigen_layers = model.module.get_fcs()
if fix_eigen_layer:
for eigen_layer in eigen_layers:
eigen_layer.eval()
for p in eigen_layer.parameters():
p.requires_grad = False
stage_name = 'restraint'
else:
model.train()
for p in model.parameters():
p.requires_grad = True
stage_name = 'relaxation'
prefix = '{}_{}_'.format(T, stage_name)
optimizer, scheduler = get_RRI_optimizer(model, lr)
for epoch in range(20):
train(epoch, model, criterion, optimizer, trainloader, use_gpu=use_gpu)
scheduler.step()
print('=> Test')
if (epoch + 1) % args.eval_freq == 0:
for name in args.target_names:
print('Evaluating {} ...'.format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
rank1 = test(model, queryloader, galleryloader, use_gpu)
save_checkpoint({
'state_dict': model.state_dict(),
'rank1': rank1,
'epoch': 0,
'arch': args.arch,
'optimizer': (),
}, args.save_dir, prefix=prefix)
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file',
type=str,
default='',
help='path to config file')
parser.add_argument(
'--gpu-devices',
type=str,
default='',
)
parser.add_argument('opts',
default=None,
nargs=argparse.REMAINDER,
help='Modify config options using the command-line')
args = parser.parse_args()
cfg = get_default_config()
if args.config_file:
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
log_name = 'test.log' if cfg.test.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
print('Show configuration\n{}\n'.format(cfg))
torch.backends.cudnn.benchmark = True
datamanager = ImageDataManager(**imagedata_kwargs(cfg))
trainloader, queryloader, galleryloader = datamanager.return_dataloaders()
print('Building model: {}'.format(cfg.model.name))
model = build_model(cfg.model.name,
datamanager.num_train_pids,
'softmax',
pretrained=cfg.model.pretrained)
if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
load_pretrained_weights(model, cfg.model.load_weights)
model = nn.DataParallel(model).cuda()
criterion = CrossEntropyLoss(datamanager.num_train_pids,
label_smooth=cfg.loss.softmax.label_smooth)
optimizer = build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = build_lr_scheduler(optimizer, **lr_scheduler_kwargs(cfg))
if cfg.model.resume and check_isfile(cfg.model.resume):
cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume,
model,
optimizer=optimizer)
if cfg.test.evaluate:
distmat = evaluate(model,
queryloader,
galleryloader,
dist_metric=cfg.test.dist_metric,
normalize_feature=cfg.test.normalize_feature,
rerank=cfg.test.rerank,
return_distmat=True)
if cfg.test.visrank:
visualize_ranked_results(distmat,
datamanager.return_testdataset(),
'image',
width=cfg.data.width,
height=cfg.data.height,
save_dir=osp.join(cfg.data.save_dir,
'visrank'))
return
time_start = time.time()
print('=> Start training')
for epoch in range(cfg.train.start_epoch, cfg.train.max_epoch):
train(epoch,
cfg.train.max_epoch,
model,
criterion,
optimizer,
trainloader,
fixbase_epoch=cfg.train.fixbase_epoch,
open_layers=cfg.train.open_layers)
scheduler.step()
if (epoch + 1) % cfg.test.eval_freq == 0 or (epoch +
1) == cfg.train.max_epoch:
rank1 = evaluate(model,
queryloader,
galleryloader,
dist_metric=cfg.test.dist_metric,
normalize_feature=cfg.test.normalize_feature,
rerank=cfg.test.rerank)
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': epoch + 1,
'rank1': rank1,
'optimizer': optimizer.state_dict(),
}, cfg.data.save_dir)
elapsed = round(time.time() - time_start)
elapsed = str(datetime.timedelta(seconds=elapsed))
print('Elapsed {}'.format(elapsed))
def main():
global args
set_random_seed(args.seed)
if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
log_name = 'log_test.txt' if args.evaluate else 'log_train.txt'
sys.stdout = Logger(osp.join(args.save_dir, log_name))
print('==========\nArgs:{}\n=========='.format(args))
if use_gpu:
print('Currently using GPU {}'.format(args.gpu_devices))
cudnn.benchmark = True
else:
warnings.warn('Currently using CPU, however, GPU is highly recommended')
print('Initializing video data manager')
dm = VideoDataManager(use_gpu, **video_dataset_kwargs(args))
trainloader, testloader_dict = dm.return_dataloaders()
print('Initializing model: {}'.format(args.arch))
model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent', 'htri'}, pretrained=not args.no_pretrained, use_gpu=use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(model)))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
model = nn.DataParallel(model).cuda() if use_gpu else model
criterion = CrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optimizer(model, **optimizer_kwargs(args))
scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs(args))
if args.resume and check_isfile(args.resume):
args.start_epoch = resume_from_checkpoint(args.resume, model, optimizer=optimizer)
if args.evaluate:
print('Evaluate only')
for name in args.target_names:
print('Evaluating {} ...'.format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
distmat = test(model, queryloader, galleryloader, args.pool_tracklet_features, use_gpu, return_distmat=True)
if args.visualize_ranks:
visualize_ranked_results(
distmat, dm.return_testdataset_by_name(name),
save_dir=osp.join(args.save_dir, 'ranked_results', name),
topk=20
)
return
time_start = time.time()
ranklogger = RankLogger(args.source_names, args.target_names)
print('=> Start training')
if args.fixbase_epoch > 0:
print('Train {} for {} epochs while keeping other layers frozen'.format(args.open_layers, args.fixbase_epoch))
initial_optim_state = optimizer.state_dict()
for epoch in range(args.fixbase_epoch):
train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu, fixbase=True)
print('Done. All layers are open to train for {} epochs'.format(args.max_epoch))
optimizer.load_state_dict(initial_optim_state)
for epoch in range(args.start_epoch, args.max_epoch):
train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_freq > 0 and (epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch:
print('=> Test')
for name in args.target_names:
print('Evaluating {} ...'.format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
rank1 = test(model, queryloader, galleryloader, args.pool_tracklet_features, use_gpu)
ranklogger.write(name, epoch + 1, rank1)
save_checkpoint({
'state_dict': model.state_dict(),
'rank1': rank1,
'epoch': epoch + 1,
'arch': args.arch,
'optimizer': optimizer.state_dict(),
}, args.save_dir)
elapsed = round(time.time() - time_start)
elapsed = str(datetime.timedelta(seconds=elapsed))
print('Elapsed {}'.format(elapsed))
ranklogger.show_summary()
