def configure_lr_scheduler(args, optimizer):
lr_scheduler = None
with logger.LoggingBlock("Learning Rate Scheduler", emph=True):
logging.info("class: %s" % args.lr_scheduler)
if args.lr_scheduler is not None:
# ----------------------------------------------
# Figure out lr_scheduler arguments
# ----------------------------------------------
kwargs = tools.kwargs_from_args(args, "lr_scheduler")
# -------------------------------------------
# Print arguments
# -------------------------------------------
for param, default in sorted(kwargs.items()):
logging.info("%s: %s" % (param, default))
# -------------------------------------------
# Add optimizer
# -------------------------------------------
kwargs["optimizer"] = optimizer
# -------------------------------------------
# Create lr_scheduler instance
# -------------------------------------------
lr_scheduler = tools.instance_from_kwargs(args.lr_scheduler_class,
kwargs)
return lr_scheduler
def configure_runtime_augmentations(args):
with logger.LoggingBlock("Runtime Augmentations", emph=True):
training_augmentation = None
validation_augmentation = None
# ----------------------------------------------------
# Training Augmentation
# ----------------------------------------------------
if args.training_augmentation is not None:
kwargs = tools.kwargs_from_args(args, "training_augmentation")
logging.info("training_augmentation: %s" %
args.training_augmentation)
for param, default in sorted(kwargs.items()):
logging.info(" %s: %s" % (param, default))
kwargs["args"] = args
training_augmentation = tools.instance_from_kwargs(
args.training_augmentation_class, kwargs)
if args.cuda:
training_augmentation = training_augmentation.cuda()
else:
logging.info("training_augmentation: None")
# ----------------------------------------------------
# Training Augmentation
# ----------------------------------------------------
if args.validation_augmentation is not None:
kwargs = tools.kwargs_from_args(args, "validation_augmentation")
logging.info("validation_augmentation: %s" %
args.training_augmentation)
for param, default in sorted(kwargs.items()):
logging.info(" %s: %s" % (param, default))
kwargs["args"] = args
validation_augmentation = tools.instance_from_kwargs(
args.validation_augmentation_class, kwargs)
if args.cuda:
validation_augmentation = validation_augmentation.cuda()
else:
logging.info("validation_augmentation: None")
return training_augmentation, validation_augmentation
def configure_optimizer(args, model_and_loss):
optimizer = None
with logger.LoggingBlock("Optimizer", emph=True):
if args.optimizer is not None:
if model_and_loss.num_parameters() == 0:
logging.info("No trainable parameters detected.")
logging.info("Setting optimizer to None.")
else:
logging.info(args.optimizer)
# -------------------------------------------
# Figure out all optimizer arguments
# -------------------------------------------
all_kwargs = tools.kwargs_from_args(args, "optimizer")
# -------------------------------------------
# Get the split of param groups
# -------------------------------------------
kwargs_without_groups = {
key: value
for key, value in all_kwargs.items() if key != "group"
}
param_groups = all_kwargs["group"]
# ----------------------------------------------------------------------
# Print arguments (without groups)
# ----------------------------------------------------------------------
for param, default in sorted(kwargs_without_groups.items()):
logging.info("%s: %s" % (param, default))
# ----------------------------------------------------------------------
# Construct actual optimizer params
# ----------------------------------------------------------------------
kwargs = dict(kwargs_without_groups)
if param_groups is None:
# ---------------------------------------------------------
# Add all trainable parameters if there is no param groups
# ---------------------------------------------------------
all_trainable_parameters = _generate_trainable_params(
model_and_loss)
kwargs["params"] = all_trainable_parameters
else:
# -------------------------------------------
# Add list of parameter groups instead
# -------------------------------------------
trainable_parameter_groups = []
dnames, dparams = _param_names_and_trainable_generator(
model_and_loss)
dnames = set(dnames)
dparams = set(list(dparams))
with logger.LoggingBlock("parameter_groups:"):
for group in param_groups:
# log group settings
group_match = group["params"]
group_args = {
key: value
for key, value in group.items()
if key != "params"
}
with logger.LoggingBlock(
"%s: %s" % (group_match, group_args)):
# retrieve parameters by matching name
gnames, gparams = _param_names_and_trainable_generator(
model_and_loss, match=group_match)
# log all names affected
for n in sorted(gnames):
logging.info(n)
# set generator for group
group_args["params"] = gparams
# append parameter group
trainable_parameter_groups.append(group_args)
# update remaining trainable parameters
dnames -= set(gnames)
dparams -= set(list(gparams))
# append default parameter group
trainable_parameter_groups.append(
{"params": list(dparams)})
# and log its parameter names
with logger.LoggingBlock("default:"):
for dname in sorted(dnames):
logging.info(dname)
# set params in optimizer kwargs
kwargs["params"] = trainable_parameter_groups
# -------------------------------------------
# Create optimizer instance
# -------------------------------------------
optimizer = tools.instance_from_kwargs(args.optimizer_class,
kwargs)
return optimizer
def configure_data_loaders(args):
with logger.LoggingBlock("Datasets", emph=True):
def _sizes_to_str(value):
if np.isscalar(value):
return '[1L]'
else:
return ' '.join([str([d for d in value.size()])])
def _log_statistics(dataset, prefix, name):
with logger.LoggingBlock("%s Dataset: %s" % (prefix, name)):
example_dict = dataset[
0] # get sizes from first dataset example
for key, value in sorted(example_dict.items()):
if key in ["index",
"basename"]: # no need to display these
continue
if isinstance(value, str):
logging.info("{}: {}".format(key, value))
else:
logging.info("%s: %s" % (key, _sizes_to_str(value)))
logging.info("num_examples: %i" % len(dataset))
# -----------------------------------------------------------------------------------------
# GPU parameters -- turning off pin_memory? for resolving the deadlock?
# -----------------------------------------------------------------------------------------
gpuargs = {
"num_workers": args.num_workers,
"pin_memory": False
} if args.cuda else {}
train_loader = None
validation_loader = None
inference_loader = None
# -----------------------------------------------------------------------------------------
# Training dataset
# -----------------------------------------------------------------------------------------
if args.training_dataset is not None:
# ----------------------------------------------
# Figure out training_dataset arguments
# ----------------------------------------------
kwargs = tools.kwargs_from_args(args, "training_dataset")
kwargs["is_cropped"] = True
kwargs["args"] = args
# ----------------------------------------------
# Create training dataset
# ----------------------------------------------
train_dataset = tools.instance_from_kwargs(
args.training_dataset_class, kwargs)
# ----------------------------------------------
# Create training loader
# ----------------------------------------------
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
drop_last=False,
**gpuargs)
_log_statistics(train_dataset,
prefix="Training",
name=args.training_dataset)
# -----------------------------------------------------------------------------------------
# Validation dataset
# -----------------------------------------------------------------------------------------
if args.validation_dataset is not None:
# ----------------------------------------------
# Figure out validation_dataset arguments
# ----------------------------------------------
kwargs = tools.kwargs_from_args(args, "validation_dataset")
kwargs["is_cropped"] = True
kwargs["args"] = args
# ----------------------------------------------
# Create validation dataset
# ----------------------------------------------
validation_dataset = tools.instance_from_kwargs(
args.validation_dataset_class, kwargs)
# ----------------------------------------------
# Create validation loader
# ----------------------------------------------
validation_loader = DataLoader(validation_dataset,
batch_size=args.batch_size_val,
shuffle=False,
drop_last=False,
**gpuargs)
_log_statistics(validation_dataset,
prefix="Validation",
name=args.validation_dataset)
return train_loader, validation_loader, inference_loader
def configure_model_and_loss(args):
# ----------------------------------------------------
# Dynamically load model and loss class with parameters
# passed in via "--model_[param]=[value]" or "--loss_[param]=[value]" arguments
# ----------------------------------------------------
with logger.LoggingBlock("Model and Loss", emph=True):
# ----------------------------------------------------
# Model
# ----------------------------------------------------
kwargs = tools.kwargs_from_args(args, "model")
kwargs["args"] = args
if type(args.checkpoint) == list and len(args.checkpoint) > 1:
models = nn.ModuleList([
tools.instance_from_kwargs(args.model_class, kwargs)
for _ in args.checkpoint
])
else:
models = tools.instance_from_kwargs(args.model_class, kwargs)
if hasattr(args, 'avoid_list') and args.avoid_list:
for model in models:
model.avoid_list = args.avoid_list.split(',')
# ----------------------------------------------------
# Training loss
# ----------------------------------------------------
training_loss = None
if args.training_loss is not None:
kwargs = tools.kwargs_from_args(args, "training_loss")
kwargs["args"] = args
training_loss = tools.instance_from_kwargs(
args.training_loss_class, kwargs)
# ----------------------------------------------------
# Validation loss
# ----------------------------------------------------
validation_loss = None
if args.validation_loss is not None:
kwargs = tools.kwargs_from_args(args, "validation_loss")
kwargs["args"] = args
validation_loss = tools.instance_from_kwargs(
args.validation_loss_class, kwargs)
# ----------------------------------------------------
# Model and loss
# ----------------------------------------------------
model_and_loss = ModelAndLoss(args, models, training_loss,
validation_loss)
# -----------------------------------------------------------
# If Cuda, transfer model to Cuda and wrap with DataParallel.
# -----------------------------------------------------------
if args.cuda:
model_and_loss = model_and_loss.cuda()
# ---------------------------------------------------------------
# Report some network statistics
# ---------------------------------------------------------------
logging.info("Batch Size: %i" % args.batch_size)
logging.info("GPGPU: Cuda") if args.cuda else logging.info(
"GPGPU: off")
logging.info("Network: %s" % args.model)
logging.info("Number of parameters: %i" %
tools.x2module(model_and_loss).num_parameters())
if training_loss is not None:
logging.info("Training Key: %s" % args.training_key)
logging.info("Training Loss: %s" % args.training_loss)
if validation_loss is not None:
logging.info("Validation Key: %s" % args.validation_key)
logging.info("Validation Loss: %s" % args.validation_loss)
return model_and_loss