task

Base physics task-specific Model class(es).

class graphnet.models.task.task.Task(*args, **kwargs)[source]

Bases: Model

Base class for all reconstruction and classification tasks.

Construct Task.

Parameters:
  • hidden_size (int) – The number of nodes in the layer feeding into this tasks, used to construct the affine transformation to the predicted quantity.

  • loss_function (LossFunction) – Loss function appropriate to the task.

  • target_labels (Union[str, List[str], None], default: None) – Name(s) of the quantity/-ies being predicted, used to extract the target tensor(s) from the Data object in .compute_loss(…).

  • prediction_labels (Union[str, List[str], None], default: None) – The name(s) of each column that is predicted by the model during inference. If not given, the name will auto matically be set to target_label + _pred.

  • transform_prediction_and_target (Optional[Callable], default: None) – Optional function to transform both the predicted and target tensor before passing them to the loss function. Useful e.g. for having the model predict quantities on a physical scale, but transforming this scale to O(1) for a numerically stable loss computation.

  • transform_target (Optional[Callable], default: None) – Optional function to transform only the target tensor before passing it, and the predicted tensor, to the loss function. Useful e.g. for having the model predict a transformed version of the target quantity, e.g. the log10- scaled energy, rather than the physical quantity itself. Used in conjunction with transform_inference to perform the inverse transform on the predicted quantity to recover the physical scale.

  • transform_inference (Optional[Callable], default: None) – Optional function to inverse-transform the model prediction to recover a physical scale. Used in conjunction with transform_target.

  • transform_support (Optional[Tuple], default: None) – Optional tuple to specify minimum and maximum of the range of validity for the inverse transforms transform_target and transform_inference in case this is restricted. By default the invertibility of transform_target is tested on the range [-1e6, 1e6].

  • loss_weight (Optional[str], default: None) – Name of the attribute in data containing per-event loss weights.

  • args (Any) –

  • kwargs (Any) –

Return type:

object

abstract property nb_inputs: int

Return number of inputs assumed by task.

abstract property default_target_labels: List[str]

Return default target labels.

abstract property default_prediction_labels: List[str]

Return default prediction labels.

forward(x)[source]

Forward pass.

Return type:

Union[Tensor, Data]

Parameters:

x (Tensor | Data) –

compute_loss(pred, data)[source]

Compute loss of pred wrt.

target labels in data.

Return type:

Tensor

Parameters:
  • pred (Tensor | Data) –

  • data (Data) –

inference()[source]

Activate inference mode.

Return type:

None

train_eval()[source]

Deactivate inference mode.

Return type:

None

class graphnet.models.task.task.IdentityTask(*args, **kwargs)[source]

Bases: Task

Identity, or trivial, task.

Construct IdentityTask.

Return the nb_outputs as a direct, affine transformation of the last hidden layer.

Parameters:
  • args (Any) –

  • kwargs (Any) –

Return type:

object

property default_target_labels: List[str]

Return default target labels.

property default_prediction_labels: List[str]

Return default prediction labels.

property nb_inputs: int

Return number of inputs assumed by task.