"""Optimization/metrics utility module.
This module contains utility functions and tools used by the other modules of this package.
"""
import copy
import inspect
import logging
import numpy as np
import torch
import thelper.optim.metrics
import thelper.utils
logger = logging.getLogger(__name__)
[docs]def create_loss_fn(config, model, loader=None, uploader=None):
"""Instantiates and returns the loss function to use for training.
The default way to specify the loss function to use is to provide a callable type to instantiate
as well as its initialization parameters. For example::
# ...
"loss": {
# if we want to use PyTorch's cross entropy loss:
# >>> loss_fn = torch.nn.CrossEntropyLoss(**params)
# >>> ...
# >>> loss = loss_fn(pred, target)
"type": "torch.nn.CrossEntropyLoss"
"params": {
"weight": [ ... ],
"reduction": "mean",
# ...
}
},
# ...
The loss function can also be queried from a member function of the model class, as such::
# ...
"loss": {
# to query the model for the loss function:
# >>> loss_fn = model.get_loss_fn(**params)
# >>> ...
# >>> loss = loss_fn(pred, target)
"model_getter": "get_loss_fn"
"params": {
# ...
}
},
# ...
If the model is supposed to compute its own loss, we suggest creating a specialized trainer class. In that
case only, the 'loss' field can be omitted from the session configuration file.
If the task is related to image classification or semantic segmentation, the classes can be weighted based
on extra parameters in the loss configuration. The strategy used to compute the weights is related to the
one in :class:`thelper.data.samplers.WeightedSubsetRandomSampler`. The exact parameters that are expected
for class reweighting are the following:
- ``weight_distribution`` (mandatory, toggle): the dictionary of weights assigned to each class, or the
rebalancing strategy to use. If omitted entirely, no class weighting will be performed.
- ``weight_param_name`` (optional, default="weight"): name of the constructor parameter that expects the
weight list.
- ``weight_max`` (optional, default=inf): the maximum weight that can be assigned to a class.
- ``weight_min`` (optional, default=0): the minimum weight that can be assigned to a class.
- ``weight_norm`` (optional, default=True): specifies whether the weights should be normalized or not.
This function also supports an extra special parameter if the task is related to semantic segmentation:
``ignore_index``. If this parameter is found and not ``None`` (integer), then the loss function will ignore
the given value when computing the loss of a sample. The exact parameters that are expected in this case are
the following:
- ``ignore_index_param_name`` (optional, default="ignore_index"): name of the constructor parameter that expects
the ignore value.
- ``ignore_index_label_name`` (optional, default="dontcare"): name of the label to pass the ignore value from.
"""
# todo: add flag to toggle loss comp in validation? (add to trainer config maybe?)
logger.debug("loading loss function")
if isinstance(model, torch.nn.DataParallel):
model = model.module # to avoid interface getter issues
def converter(x):
if isinstance(x, (list, np.ndarray)) and all([np.isscalar(i) and not isinstance(i, str) for i in x]):
x = torch.FloatTensor(x)
return x if uploader is None else uploader(x)
if not isinstance(config, dict):
raise AssertionError("config should be provided as a dictionary")
if "model_getter" in config and "type" in config:
raise AssertionError("loss config cannot have both 'model_attrib_name' and 'type' fields")
if "model_getter" in config:
model_getter_name = config["model_getter"]
if not isinstance(model_getter_name, str):
raise AssertionError("unexpected model getter name type")
if not hasattr(model, model_getter_name) or not callable(getattr(model, model_getter_name)):
raise AssertionError("invalid model getter attribute")
loss_type = getattr(model, model_getter_name)
elif "type" in config:
loss_type = thelper.utils.import_class(config["type"])
else:
raise AssertionError("loss config missing 'type' or 'model_attrib_name' field")
loss_params = thelper.utils.get_key_def(["params", "parameters"], config, {})
loss_params = copy.deepcopy(loss_params) # required here, we might add some parameters below
weight_param_name = thelper.utils.get_key_def("weight_param_name", config, "weight")
if thelper.utils.str2bool(thelper.utils.get_key_def("weight_classes", config, False)) or \
thelper.utils.get_key_def("weight_distribution", config, None) is not None:
if not thelper.utils.str2bool(thelper.utils.get_key_def("weight_classes", config, True)):
raise AssertionError("'weight_classes' now deprecated, set 'weight_distribution' directly to toggle on")
if not isinstance(model.task, (thelper.tasks.Classification, thelper.tasks.Segmentation)):
raise AssertionError("task type does not support class weighting")
weight_distrib = thelper.utils.get_key("weight_distribution", config,
msg="missing 'weight_distribution' field in loss config")
if isinstance(weight_distrib, dict):
for label, weight in weight_distrib.items():
if label not in model.task.class_names:
raise AssertionError("weight distribution label '%s' not in dataset class list" % label)
if not isinstance(weight, float):
raise AssertionError("expected weight distrib map to provide weights as floats directly")
elif isinstance(weight_distrib, str):
weight_max = float("inf")
if "weight_max" in config:
weight_max = float(config["weight_max"])
weight_min = 0
if "weight_min" in config:
weight_min = float(config["weight_min"])
if weight_distrib != "uniform":
if loader is None or not loader:
raise AssertionError("cannot get class sizes, no training data available")
label_sizes_map = model.task.get_class_sizes(loader.dataset)
weight_norm = False
else:
label_sizes_map = {label: -1 for label in model.task.class_names} # counts don't matter
weight_norm = True
if "weight_norm" in config:
weight_norm = thelper.utils.str2bool(config["weight_norm"])
weight_distrib = thelper.data.utils.get_class_weights(label_sizes_map, weight_distrib, invmax=True,
maxw=weight_max, minw=weight_min, norm=weight_norm)
else:
raise AssertionError("unexpected weight distribution strategy (should be map or string)")
weight_list_str = "weight_distribution: {"
for label, weight in weight_distrib.items():
weight_list_str += "\n \"%s\": %s," % (label, weight)
logger.info(weight_list_str + "\n}")
loss_params[weight_param_name] = [weight_distrib[label] if label in weight_distrib
else 1.0 for label in model.task.class_names]
if isinstance(model.task, thelper.tasks.Segmentation):
ignore_index_param_name = thelper.utils.get_key_def("ignore_index_param_name", config, "ignore_index")
ignore_index_label_name = thelper.utils.get_key_def("ignore_index_label_name", config, "dontcare")
loss_sig = inspect.signature(loss_type)
if ignore_index_param_name in loss_sig.parameters:
if ignore_index_label_name != "dontcare":
loss_params[ignore_index_param_name] = model.task.class_indices[ignore_index_label_name]
else:
loss_params[ignore_index_param_name] = model.task.dontcare
if loss_params[ignore_index_param_name] is None:
# some loss functions dont actually accept 'None' as the dontcare value (e.g. cross entropy loss)
# ... switch back to the default value in that case
loss_params[ignore_index_param_name] = loss_sig.parameters[ignore_index_param_name].default
if weight_param_name in loss_params:
loss_params[weight_param_name] = converter(loss_params[weight_param_name])
loss = loss_type(**loss_params)
return loss
[docs]def create_metrics(config):
"""Instantiates and returns the metrics defined in the configuration dictionary.
All arguments are expected to be handed in through the configuration via a dictionary named 'params'.
"""
return thelper.train.utils.create_consumers(config)
[docs]def create_optimizer(config, model):
"""Instantiates and returns the optimizer to use for training.
By default, the optimizer will be instantiated with the model parameters given as the first argument
of its constructor. All supplementary arguments are expected to be handed in through the configuration
via a dictionary named 'params'.
"""
logger.debug("loading optimizer")
if not isinstance(config, dict):
raise AssertionError("config should be provided as a dictionary")
if "type" not in config or not config["type"]:
raise AssertionError("optimizer config missing 'type' field")
optimizer_type = thelper.utils.import_class(config["type"])
optimizer_params = thelper.utils.get_key_def(["params", "parameters"], config, {})
optimizer = optimizer_type(filter(lambda p: p.requires_grad, model.parameters()), **optimizer_params)
return optimizer
[docs]def create_scheduler(config, optimizer):
"""Instantiates and returns the learning rate scheduler to use for training.
All arguments are expected to be handed in through the configuration via a dictionary named 'params'.
"""
logger.debug("loading scheduler")
if not isinstance(config, dict):
raise AssertionError("config should be provided as a dictionary")
if "type" not in config or not config["type"]:
raise AssertionError("scheduler config missing 'type' field")
scheduler_type = thelper.utils.import_class(config["type"])
scheduler_params = thelper.utils.get_key_def(["params", "parameters"], config, {})
scheduler = scheduler_type(optimizer, **scheduler_params)
scheduler_step_metric = None
if "step_metric" in config:
scheduler_step_metric = config["step_metric"]
return scheduler, scheduler_step_metric
[docs]def get_lr(optimizer):
"""Returns the optimizer's learning rate, or 0 if not found."""
if optimizer is not None:
for param_group in optimizer.param_groups:
if "lr" in param_group:
return param_group["lr"]
return 0
