def init_model(self):
print('Initializing model: {}'.format(args.model))
self.model = models.init_model(name=self.args.model,
num_classes=self.dm.num_attributes,
pretrained=not self.args.no_pretrained,
use_gpu=self.use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(self.model)))
# Load pretrained weights if specified in args.
self.loaded_args = self.args
load_file = osp.join(args.save_experiment, args.load_weights)
if args.load_weights:
# TODO: implement result dict
if check_isfile(load_file):
cp = load_pretrained_weights([self.model], load_file)
if "args" in cp:
self.loaded_args = cp["args"]
else:
print("WARNING: Could not load args. ")
else:
print("WARNING: Could not load pretrained weights")
# Load model onto GPU if GPU is used.
self.model = self.model.cuda() if self.use_gpu else self.model
# Select Loss function.
if args.loss_func == "deepmar":
pos_ratio = self.dm.dataset.get_positive_attribute_ratio()
self.criterion = DeepMARLoss(pos_ratio,
args.train_batch_size,
use_gpu=self.use_gpu,
sigma=args.loss_func_param)
elif args.loss_func == "scel":
self.criterion = SigmoidCrossEntropyLoss(
num_classes=self.dm.num_attributes, use_gpu=self.use_gpu)
elif args.loss_func == "sscel":
attribute_grouping = self.dm.dataset.attribute_grouping
self.criterion = SplitSoftmaxCrossEntropyLoss(attribute_grouping,
use_gpu=self.use_gpu)
else:
self.criterion = None
self.f1_calibration_thresholds = None
self.optimizer = init_optimizer(self.model, **optimizer_kwargs(args))
self.scheduler = init_lr_scheduler(self.optimizer,
**lr_scheduler_kwargs(args))
self.model = nn.DataParallel(
self.model) if self.use_gpu else self.model
# Set default for max_epoch if it was not passed as an argument in the console.
if self.args.max_epoch < 0:
self.args.max_epoch = 30
self.model_list = [self.model]
self.optimizer_list = [self.optimizer]
self.scheduler_list = [self.scheduler]
self.criterion_list = [self.criterion]
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:
print('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)
if args.resume and check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch'] + 1
best_rank1 = checkpoint['rank1']
print('Loaded checkpoint from "{}"'.format(args.resume))
print('- start_epoch: {}\n- rank1: {}'.format(args.start_epoch,
best_rank1))
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.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
start_time = time.time()
ranklogger = RankLogger(args.source_names, args.target_names)
train_time = 0
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):
start_train_time = time.time()
train(epoch,
model,
criterion_xent,
criterion_htri,
optimizer,
trainloader,
use_gpu,
fixbase=True)
train_time += round(time.time() - start_train_time)
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):
start_train_time = time.time()
train(epoch, model, criterion_xent, criterion_htri, optimizer,
trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
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,
}, False,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
'Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.'.
format(elapsed, train_time))
ranklogger.show_summary()
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 image data manager')
dm = ImageDataManager(use_gpu, **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)
model = models.init_model(name=args.arch,
num_classes=dm.num_train_pids,
image_size=128,
patch_size=8,
dim=1024,
depth=6,
heads=16,
mlp_dim=2048,
loss={'xent', 'htri'},
pretrained=not args.no_pretrained)
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
# print('Initializing model: ViT')
# model = ViT(
# image_size = 128,
# patch_size = 8,
# num_classes = dm.num_train_pids,
# dim = 1024,
# depth = 6,
# heads = 16,
# mlp_dim = 2048,
# channels = 3,
# dropout = 0.1,
# emb_dropout = 0.1,
# loss = 'xent'
# )
# model = nn.DataParallel(model).cuda() if use_gpu else model
criterion_xent = 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,
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, 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()
def init_model(self):
print('Initializing main model: {}'.format(args.model))
self.model_main = models.init_model(
name=self.args.model,
num_classes=self.dm.num_attributes,
pretrained=not self.args.no_pretrained,
use_gpu=self.use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(self.model_main)))
print('Initializing HP model: {}'.format(args.hp_model))
# Determine the size of the output vector for the HP-Net.
num_hp_net_outputs = 1 if self.args.hp_net_simple else self.dm.num_attributes
# Init the HP-Net
self.model_hp = models.init_model(
name="hp_net_" + self.args.hp_model,
num_classes=num_hp_net_outputs,
pretrained=not self.args.no_pretrained)
print('Model size: {:.3f} M'.format(count_num_param(self.model_hp)))
if self.args.rejector == "none":
self.rejector = rejectors.NoneRejector()
elif self.args.rejector == 'macc':
self.rejector = rejectors.MeanAccuracyRejector(
self.args.max_rejection_quantile)
elif self.args.rejector == "median":
self.rejector = rejectors.MedianRejector(
self.args.max_rejection_quantile)
elif self.args.rejector == "threshold":
self.rejector = rejectors.ThresholdRejector(
self.args.rejection_threshold,
self.args.max_rejection_quantile)
elif self.args.rejector == "quantile":
self.rejector = rejectors.QuantileRejector(
self.args.max_rejection_quantile)
elif self.args.rejector == 'f1':
self.rejector = rejectors.F1Rejector(
self.args.max_rejection_quantile)
else:
raise ValueError("Unsupported rejection strategy: '{}'".format(
self.args.rejector))
print("Using rejection strategy '{}'".format(self.args.rejector))
if self.args.hp_calib == 'none':
self.hp_calibrator = NoneCalibrator()
elif self.args.hp_calib == 'linear':
self.hp_calibrator = LinearCalibrator()
else:
raise ValueError("Unsupported calibrator: '{}'".format(
self.args.hp_calib))
print("Using calibrator for HP-Loss '{}'".format(self.args.hp_calib))
# Load pretrained weights if specified in args.
load_file = osp.join(args.save_experiment, args.load_weights)
self.loaded_args = self.args
if args.load_weights:
if check_isfile(load_file):
cp = load_pretrained_weights([self.model_main, self.model_hp],
load_file)
if "args" in cp:
self.loaded_args = cp["args"]
else:
print("WARNING: Could not load args. ")
if "result_dict" in cp and cp[
"result_dict"] is not None and self.args.evaluate:
self.result_dict = cp["result_dict"]
self.result_manager = ResultManager(self.result_dict)
print("Loaded result dict with keys: ")
print(sorted(list(self.result_dict.keys())))
if "rejection_thresholds" in self.result_dict:
self.rejector.load_thresholds(
self.result_dict["rejection_thresholds"])
if self.rejector.is_initialized():
print("Loaded rejection thresholds. ")
else:
print(
"Loaded uninitialized (None) rejection thresholds. "
)
else:
print("WARNING: Could not load rejection thresholds. ")
else:
print("WARNING: Could not load pretrained weights")
self.new_eval_split = self.args.eval_split != self.loaded_args.eval_split
# Load model onto GPU if GPU is used.
self.model_main = self.model_main.cuda(
) if self.use_gpu else self.model_main
self.model = self.model_main
self.model_hp = self.model_hp.cuda() if self.use_gpu else self.model_hp
self.pos_ratio = self.dm.dataset.get_positive_attribute_ratio()
# Select Loss function.
# Select Loss function.
if args.loss_func == "deepmar":
self.criterion = DeepMARLoss(self.pos_ratio,
args.train_batch_size,
use_gpu=self.use_gpu,
sigma=args.loss_func_param)
elif args.loss_func == "scel":
self.criterion = SigmoidCrossEntropyLoss(
num_classes=self.dm.num_attributes, use_gpu=self.use_gpu)
else:
self.criterion = None
self.criterion_main = self.criterion
self.criterion_hp = HardnessPredictorLoss(
self.args.use_deepmar_for_hp,
self.pos_ratio,
self.dm.num_attributes,
use_gpu=self.use_gpu,
sigma=self.args.hp_loss_param,
use_visibility=self.args.use_bbs_gt,
visibility_weight=self.args.hp_visibility_weight)
self.f1_calibration_thresholds = None
self.optimizer_main = init_optimizer(self.model_main,
**optimizer_kwargs(args))
self.scheduler_main = init_lr_scheduler(self.optimizer_main,
**lr_scheduler_kwargs(args))
self.optimizer = self.optimizer_main
self.scheduler = self.scheduler_main
op_args = optimizer_kwargs(args)
sc_args = lr_scheduler_kwargs(args)
op_args['lr'] *= self.args.hp_net_lr_multiplier
self.optimizer_hp = init_optimizer(self.model_hp, **op_args)
sc_args["stepsize"] = [
i + self.args.hp_epoch_offset for i in sc_args["stepsize"]
]
self.scheduler_hp = init_lr_scheduler(self.optimizer_hp, **sc_args)
self.model_main = nn.DataParallel(
self.model_main) if self.use_gpu else self.model_main
self.model = self.model_main
self.model_hp = nn.DataParallel(
self.model_hp) if self.use_gpu else self.model_hp
if not self.args.evaluate:
self.init_epochs()
self.model_list = [self.model_main, self.model_hp]
self.optimizer_list = [self.optimizer_main, self.optimizer_hp]
self.scheduler_list = [self.scheduler_main, self.scheduler_hp]
self.criterion_list = [self.criterion_main, self.criterion_hp]