Task

class flash.core.model.Task(model=None, loss_fn=None, learning_rate=None, optimizer='Adam', lr_scheduler=None, metrics=None, output_transform=None)

A general Task.

Parameters

• model (Optional[TypeVar(MODEL_TYPE, Module, None)]) – Model to use for the task.

• loss_fn (Optional[TypeVar(LOSS_FN_TYPE, Callable, Mapping, Sequence, None)]) – Loss function for training.

• learning_rate (Optional[float]) – Learning rate to use for training. If None (the default) then the default LR for your chosen optimizer will be used.

• optimizer (TypeVar(OPTIMIZER_TYPE, str, Callable, Tuple[str, Dict[str, Any]], None)) – Optimizer to use for training.

• lr_scheduler (Optional[TypeVar(LR_SCHEDULER_TYPE, str, Callable, Tuple[str, Dict[str, Any]], Tuple[str, Dict[str, Any], Dict[str, Any]], None)]) – The LR scheduler to use during training.

• metrics (Optional[TypeVar(METRICS_TYPE, Metric, Mapping, Sequence, None)]) – Metrics to compute for training and evaluation. Can either be an metric from the torchmetrics package, a custom metric inheriting from torchmetrics.Metric, a callable function or a list/dict containing a combination of the aforementioned. In all cases, each metric needs to have the signature metric(preds,target) and return a single scalar tensor.

• output_transform (Optional[TypeVar(OUTPUT_TRANSFORM_TYPE, flash.core.data.io.output_transform.OutputTransform, None)]) – OutputTransform to use as the default for this task.

>>> Task()  
Task(
  (train_metrics): ModuleDict()
  (val_metrics): ModuleDict()
  (test_metrics): ModuleDict()
)

static apply_filtering(y, y_hat)

This function is used to filter some labels or predictions which aren’t conform.

Return type

Tuple[Tensor, Tensor]

as_embedder(layer)

Convert this task to an embedder. Note that the parameters are not copied so that any optimization of the embedder will also apply to the converted Task.

Parameters

layer (str) – The layer to embed to. This should be one of the available_layers().

classmethod available_finetuning_strategies()

Returns a list containing the keys of the available Finetuning Strategies.

Return type

List[str]

available_layers()

Get the list of available layers for use with the as_embedder() method.

classmethod available_lr_schedulers()

Returns a list containing the keys of the available LR schedulers.

Return type

List[str]

classmethod available_optimizers()

Returns a list containing the keys of the available Optimizers.

Return type

List[str]

classmethod available_outputs()

Returns the list of available outputs (that can be used during prediction or serving) for this Task.

Examples

..testsetup:

>>> from flash import Task

>>> print(Task.available_outputs())
['preds', 'raw']

Return type

List[str]

configure_optimizers()

Implement how optimizer and optionally learning rate schedulers should be configured.

Return type

Union[Optimizer, Tuple[List[Optimizer], List[_LRScheduler]]]

get_num_training_steps()

Total training steps inferred from datamodule and devices.

Return type

int

modules_to_freeze()

By default, we try to get the backbone attribute from the task and return it or None if not present.

Return type

Optional[Module]

Returns

The backbone Module to freeze or None if this task does not have a backbone attribute.

serve(host='127.0.0.1', port=8000, sanity_check=True, input_cls=None, transform=<class 'flash.core.data.io.input_transform.InputTransform'>, transform_kwargs=None, output=None)

Serve the Task. Override this method to provide a default input_cls, transform, and transform_kwargs.

Parameters

• host (str) – The IP address to host the Task on.

• port (int) – The port to host on.

• sanity_check (bool) – If True, runs a sanity check before serving.

• input_cls (Optional[Type[ServeInput]]) – The ServeInput type to use.

• transform (TypeVar(INPUT_TRANSFORM_TYPE, Type[flash.core.data.io.input_transform.InputTransform], Callable, Tuple[Union[StrEnum, str], Dict[str, Any]], Union[StrEnum, str], None)) – The transform to use when serving.

• transform_kwargs (Optional[Dict]) – Keyword arguments used to instantiate the transform.

Return type

Composition

step(batch, batch_idx, metrics)

Implement the core logic for the training/validation/test step. By default this includes:

• Inference on the current batch

• Calculating the loss

• Calculating relevant metrics

Override for custom behavior.

Parameters

• batch (Any) – The output of your dataloader. Can either be a single Tensor or a list of Tensors.

• batch_idx (int) – Integer displaying index of this batch

• metrics (ModuleDict) – A module dict containing metrics for calculating relevant training statitics

Return type

Any

Returns

A dict containing both the loss and relevant metrics