Task Heads

salt.models.TaskBase

Bases: torch.nn.Module, abc.ABC

Base class for task heads.

Tasks wrap a dense network, a loss, a target label, and a scalar weight.

Parameters:

Name	Type	Description	Default
name	str	Arbitrary name of the task, used for logging and inference.	required
input_name	str	Name of the input stream consumed by this task (e.g., "jet", "track", "objects").	required
dense_config	dict	Keyword arguments for :class:salt.models.Dense, the head producing the task outputs.	required
loss	torch.nn.Module	Loss function applied to the head outputs.	required
weight	float	Scalar multiplier for the task loss in the overall objective. The default is 1.0.	1.0

Source code in salt/models/task.py

def __init__(
    self,
    name: str,
    input_name: str,
    dense_config: dict,
    loss: nn.Module,
    weight: float = 1.0,
):
    super().__init__()

    self.name = name
    self.input_name = input_name
    self.net = Dense(**dense_config)
    self.loss = loss
    self.weight = weight

salt.models.ClassificationTask

Bases: salt.models.task.TaskBase

Multi-class or binary classification task head.

Parameters:

Name	Type	Description	Default
label	str	Label name for the task.	required
class_names	list[str] | None	Ordered class names (index-aligned with outputs). If None, attempt to infer from the label via :data:CLASS_NAMES.	None
label_map	collections.abc.Mapping | None	Mapping to remap integer labels for training (e.g., {0,4,5} → {0,1,2}).	None
sample_weight	str | None	Key of a per-sample weight found in labels_dict[self.input_name]. Requires the configured loss to have reduction="none".	None
use_class_dict	bool	If True, read class weights for the loss from a class dictionary.	False
**kwargs		Forwarded to :class:TaskBase.	{}

Raises:

Type	Description
ValueError	If a label map is provided without class names, or if the number of outputs does not match the number of classes.

Source code in salt/models/task.py

def __init__(
    self,
    label: str,
    class_names: list[str] | None = None,
    label_map: Mapping | None = None,
    sample_weight: str | None = None,
    use_class_dict: bool = False,
    **kwargs,
):
    super().__init__(**kwargs)
    self.label = label
    self.class_names = class_names
    self.label_map = label_map
    if self.label_map is not None and self.class_names is None:
        raise ValueError("Specify class names when using label_map.")
    if hasattr(self.loss, "ignore_index"):
        self.loss.ignore_index = -1
    self.sample_weight = sample_weight
    if self.sample_weight is not None:
        assert (
            self.loss.reduction == "none"
        ), "Sample weights only supported for reduction='none'"
    if self.class_names is None:
        self.class_names = CLASS_NAMES[self.label]
    if len(self.class_names) != self.net.output_size:
        raise ValueError(
            f"{self.name}: "
            f"Number of outputs ({self.net.output_size}) does not match "
            f"number of class names ({len(self.class_names)}). Class names: {self.class_names}"
        )
    self.use_class_dict = use_class_dict

salt.models.RegressionTaskBase

Bases: salt.models.task.TaskBase, abc.ABC

Base class for regression tasks with optional target scaling.

Parameters:

Name	Type	Description	Default
targets	list[str] | str	Regression target name(s).	required
scaler	salt.utils.scalers.RegressionTargetScaler | None	Functional scaler for targets. Mutually exclusive with other scaling options.	None
target_denominators	list[str] | str | None	Denominator variable(s) for forming ratios as targets. Mutually exclusive with other scaling options.	None
norm_params	dict | None	Mean/std normalization parameters for each target. Mutually exclusive with other scaling options. Expected keys: "mean", "std".	None
custom_output_names	list[str] | str | None	Optional custom output names overriding the default.	None
**kwargs		Forwarded to :class:TaskBase.	{}

Raises:

Type	Description
ValueError	If multiple scaling methods are set simultaneously or if parameter counts do not match the number of targets.

Source code in salt/models/task.py

def __init__(
    self,
    targets: list[str] | str,
    scaler: RegressionTargetScaler | None = None,
    target_denominators: list[str] | str | None = None,
    norm_params: dict | None = None,
    custom_output_names: list[str] | str | None = None,
    **kwargs,
):
    super().__init__(**kwargs)
    self.scaler = scaler
    self.targets = listify(targets)
    self.target_denominators = listify(target_denominators)
    self.custom_output_names = listify(custom_output_names)
    if norm_params:
        norm_params["mean"] = listify(norm_params["mean"])
        norm_params["std"] = listify(norm_params["std"])
    self.norm_params = norm_params

    if [scaler, target_denominators, norm_params].count(None) not in {2, 3}:
        raise ValueError("Can only use a single scaling method")

    if self.scaler:
        for target in self.targets:
            self.scaler.scale(target, torch.Tensor(1))
    if self.target_denominators and len(self.targets) != len(self.target_denominators):
        raise ValueError(
            f"{self.name}: "
            f"Number of targets ({len(self.targets)}) does not match "
            f"number of target denominators ({len(self.target_denominators)})"
        )
    if self.norm_params and len(self.norm_params["mean"]) != len(self.targets):
        raise ValueError(
            f"{self.name}: "
            f"Number of means in norm_params ({len(self.norm_params['mean'])}) does not match "
            f"number of targets ({len(self.targets)})"
        )
    if self.norm_params and len(self.norm_params["std"]) != len(self.targets):
        raise ValueError(
            f"{self.name}: "
            f"Number of stds in norm_params ({len(self.norm_params['std'])}) does not match "
            f"number of targets ({len(self.targets)})"
        )

salt.models.RegressionTask

Bases: salt.models.task.RegressionTaskBase

Standard regression task head.

Parameters:

Name	Type	Description	Default
scaler	salt.utils.scalers.RegressionTargetScaler | None	Backward-compatibility placeholder; if provided, stored on the instance.	None
**kwargs		Forwarded to :class:RegressionTaskBase.	{}

Raises:

Type	Description
ValueError	If the number of outputs does not match the number of regression targets.

Source code in salt/models/task.py

def __init__(self, scaler: RegressionTargetScaler | None = None, **kwargs):
    super().__init__(**kwargs)
    if self.net.output_size != len(self.targets):
        raise ValueError(
            f"{self.name}: "
            f"Number of outputs ({self.net.output_size}) does not match "
            f"number of targets ({len(self.targets)})"
        )
    self.scaler = scaler

salt.models.GaussianRegressionTask

Bases: salt.models.task.RegressionTaskBase

Gaussian regression task head (predicts mean and variance).

The head outputs 2 * len(targets) values per example (means and variances). The loss is the negative log-likelihood under a Gaussian.

Parameters:

Name	Type	Description	Default
**kwargs	typing.Any	Forwarded to :class:RegressionTaskBase.	{}

Raises:

Type	Description
ValueError	If the number of outputs is not 2 * len(targets).

Source code in salt/models/task.py

def __init__(self, **kwargs: Any):
    super().__init__(**kwargs)
    if self.net.output_size != 2 * len(self.targets):
        raise ValueError(
            f"{self.name}: "
            f"Number of targets ({len(self.targets)}) is not twice the "
            f"number of outputs ({self.net.output_size})"
        )

---

Last update: November 16, 2023
Created: October 20, 2023