def __init__(self, seq, buffer, dataset_path, model_path, set_type='train',
max_res=800, branch_arch='alexnet', ctx_mode='max'):
"""
Args:
seq: (int) The number of the sequence according to the get_sequence
function, which mirrors the indexing of the ImageNetVID class.
buffer: (queue.Queue) The data buffer between the producerThread and
the consumer application (the display). The elements stored in
this buffer are defined by the BufferElement namedtuple.
dataset_path: (string) The path to the root of the ImageNet dataset.
model_path: (string) The path to the models .pth.tar file containing
the model's weights.
set_type: (string) The subset of the ImageNet VID dataset, can be
'train' or 'val'.
max_res: (int) The maximum resolution in pixels. If any dimension
of the image exceeds this value, the final image published by
the producer is resized (keeping the aspect ratio). Used to
balance the load between the consumer (main) thread and the
producer.
branch_arch: (string) The architecture of the branch of the siamese
net. Might be: 'alexnet', 'vgg11_5c'.
ctx_mode: (string) The strategy used to define the context region
around the target, using the bounding box dimensions. The 'max'
mode uses the biggest dimension, while the 'mean' mode uses the
mean of the dimensions.
"""
super(ProducerThread, self).__init__(daemon=True)
self.frames, self.bboxes_norm, self.valid_frames, self.vid_dims = (
get_sequence(seq, dataset_path, set_type=set_type))
self.idx = 0
self.seq_size = len(self.frames)
self.buffer = buffer
# TODO put the model info inside the checkpoint file.
if branch_arch == 'alexnet':
self.net = mdl.SiameseNet(mdl.BaselineEmbeddingNet(), stride=4)
elif branch_arch == 'vgg11_5c':
self.net = mdl.SiameseNet(mdl.VGG11EmbeddingNet_5c(), stride=4)
elif branch_arch == "vgg16_8c":
self.net = mdl.SiameseNet(mdl.VGG16EmbeddingNet_8c(), stride=4)
checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
self.net.load_state_dict(checkpoint['state_dict'])
# Tuple of (H, w), the dimensions to which the image will be resized.
self.resize_dims = None
self.net = self.net.to(device)
self.net.eval()
self.ref, self.ref_emb = self.make_ref(ctx_mode=ctx_mode)
def main(args):
""" Execute this file as a script and it will save the baseline model in
the default experiment folder.
"""
net_path = args.net_file_path
net = mdl.SiameseNet(mdl.BaselineEmbeddingNet())
net = load_baseline(net_path, net)
optimizer = optim.Adam(net.parameters(), lr=1e-3)
export_to_checkpoint({'epoch': 0,
'state_dict': net.state_dict(),
'optim_dict': optimizer.state_dict()},
checkpoint=args.dst_path)
def main(args):
root_dir = dirname(abspath(__file__))
# Load the parameters from json file
imagenet_dir = args.data_dir
exp_dir = join(root_dir, 'training', 'experiments', args.exp_name)
json_path = join(exp_dir, 'parameters.json')
assert isfile(json_path), (
"No json configuration file found at {}".format(json_path))
params = train_utils.Params(json_path)
# Add the timer option to the parameters
params.update_with_dict({'timer': args.timer})
params.update_with_dict({'num_workers': args.num_workers})
train_utils.set_logger(join(exp_dir, '{}.log'.format(args.mode)))
logging.info("----Starting train script in mode: {}----".format(args.mode))
setup_timer = Timer(convert=True)
setup_timer.reset()
logging.info("Loading datasets...")
# Get the correct model
if params.model == 'BaselineEmbeddingNet':
model = mdl.SiameseNet(mdl.BaselineEmbeddingNet(),
upscale=params.upscale,
corr_map_size=33,
stride=4)
elif params.model == 'VGG11EmbeddingNet_5c':
model = mdl.SiameseNet(mdl.VGG11EmbeddingNet_5c(),
upscale=params.upscale,
corr_map_size=33,
stride=4)
elif params.model == 'VGG16EmbeddingNet_8c':
model = mdl.SiameseNet(mdl.VGG16EmbeddingNet_8c(),
upscale=params.upscale,
corr_map_size=33,
stride=4)
# Freeze all the indicated parameters
for i, (name, parameter) in enumerate(model.named_parameters()):
if i in params.parameter_freeze:
logging.info("Freezing parameter {}".format(name))
parameter.requires_grad = False
model = model.to(device)
# Set the tensorboard summary maker
summ_maker = SummaryMaker(join(exp_dir, 'tensorboard'), params,
model.upscale_factor)
label_function = create_BCELogit_loss_label
img_read_fcn = imutils.get_decode_jpeg_fcn(flag=args.imutils_flag)
img_resize_fcn = imutils.get_resize_fcn(flag=args.imutils_flag)
logging.info("Validation dataset...")
metadata_val_file = join(exp_dir, "metadata.val")
val_set = ImageNetVID_val(imagenet_dir,
label_fcn=label_function,
pos_thr=params.pos_thr,
neg_thr=params.neg_thr,
upscale_factor=model.upscale_factor,
cxt_margin=params.context_margin,
reference_size=params.reference_sz,
search_size=params.search_sz,
img_read_fcn=img_read_fcn,
resize_fcn=img_resize_fcn,
metadata_file=metadata_val_file,
save_metadata=metadata_val_file,
max_frame_sep=params.max_frame_sep)
val_loader = DataLoader(val_set,
batch_size=params.batch_size,
shuffle=False,
num_workers=params.num_workers,
pin_memory=True)
if params.eval_epoch_size > len(val_loader):
logging.info('The user set eval_epoch_size ({}) is bigger than the '
'size of the eval set ({}). \n Setting '
'eval_epoch_size to the eval set size.'.format(
params.eval_epoch_size, len(val_loader)))
params.eval_epoch_size = len(val_loader)
# Define the model and optimizer
# fetch loss function and metrics
loss_fn = losses.BCELogit_Loss
metrics = met.METRICS
# Set the optional keyword arguments for the functions that need it
metrics['center_error']['kwargs']['upscale_factor'] = model.upscale_factor
try:
if args.mode == 'train':
logging.info("Training dataset...")
metadata_train_file = join(exp_dir, "metadata.train")
train_set = ImageNetVID(imagenet_dir,
label_fcn=label_function,
pos_thr=params.pos_thr,
neg_thr=params.neg_thr,
upscale_factor=model.upscale_factor,
cxt_margin=params.context_margin,
reference_size=params.reference_sz,
search_size=params.search_sz,
img_read_fcn=img_read_fcn,
resize_fcn=img_resize_fcn,
metadata_file=metadata_train_file,
save_metadata=metadata_train_file,
max_frame_sep=params.max_frame_sep)
train_loader = DataLoader(train_set,
batch_size=params.batch_size,
shuffle=True,
num_workers=params.num_workers,
pin_memory=True)
# Though I'm not a big fan of changing the value of a parameter
# variable after it has been read, at least I let the user know I'm
# changing it.
if params.train_epoch_size > len(train_loader):
logging.info(
'The user set train_epoch_size ({}) is bigger than the '
'size of the train set ({}). \n Setting '
'train_epoch_size to the train set size.'.format(
params.train_epoch_size, len(train_loader)))
params.train_epoch_size = len(train_loader)
logging.info("Done")
logging.info("Setup time: {}".format(setup_timer.elapsed))
parameters = filter(lambda p: p.requires_grad, model.parameters())
optimizer = optimz.OPTIMIZERS[params.optim](parameters,
**params.optim_kwargs)
# Set the scheduler, that updates the learning rate using a exponential
# decay. If you don't want lr decay set it to 1.
logging.info("Using Exponential Learning Rate Decay of {}".format(
params.lr_decay))
scheduler = torch.optim.lr_scheduler.ExponentialLR(
optimizer, params.lr_decay)
logging.info("Epoch sizes: {} in train and {} in eval".format(
params.train_epoch_size, params.eval_epoch_size))
logging.info("Starting training for {} epoch(s)".format(
params.num_epochs))
with Timer(convert=True) as t:
train_and_evaluate(model,
train_loader,
val_loader,
optimizer,
scheduler,
loss_fn,
metrics,
params,
exp_dir,
args,
summ_maker=summ_maker)
if params.timer:
logging.info(
"[profiling] Total time to train {} epochs, with {}"
" elements on training dataset and {} "
"on val dataset: {}".format(params.num_epochs,
len(train_loader),
len(val_loader), t.elapsed))
elif args.mode == 'eval':
logging.info("Done")
with Timer(convert=True) as total:
logging.info("Starting evaluation")
# TODO write a decent Exception
if args.restore_file is None:
raise IncompleteArgument("In eval mode you have to specify"
" a model checkpoint to be loaded"
" and evaluated."
" E.g: --restore_file best")
checkpoint_path = join(exp_dir, args.restore_file + '.pth.tar')
train_utils.load_checkpoint(checkpoint_path, model)
# Evaluate
summ_maker.epoch = 0
test_metrics = evaluate(model,
loss_fn,
val_loader,
metrics,
params,
args,
summ_maker=summ_maker)
save_path = join(
exp_dir, "metrics_test_{}.json".format(args.restore_file))
train_utils.save_dict_to_json(test_metrics, save_path)
if params.timer:
logging.info("[profiling] Total evaluation time: {}".format(
total.elapsed))
except KeyboardInterrupt:
logging.info("=== User interrupted execution ===")
raise
except Exception as e:
logging.exception("Fatal error in main loop")
logging.info("=== Execution Terminated with error ===")
else:
logging.info("=== Execution exited normally ===")