topobench.run module#

Main entry point for training and testing models.

class Any(*args, **kwargs)#

Bases: object

Special type indicating an unconstrained type.

Any is compatible with every type.
Any assumed to have all methods.
All values assumed to be instances of Any.

Note that all the above statements are true from the point of view of static type checkers. At runtime, Any should not be used with instance checks.

class Callback#

Bases: object

Abstract base class used to build new callbacks.

Subclass this class and override any of the relevant hooks

load_state_dict(state_dict)#

Called when loading a checkpoint, implement to reload callback state given callback’s state_dict.

Parameters:: state_dict (Dict[str, Any]) – the callback state returned by state_dict.

on_after_backward(trainer, pl_module)#

Called after loss.backward() and before optimizers are stepped.

on_before_backward(trainer, pl_module, loss)#

Called before loss.backward().

on_before_optimizer_step(trainer, pl_module, optimizer)#

Called before optimizer.step().

on_before_zero_grad(trainer, pl_module, optimizer)#

Called before optimizer.zero_grad().

on_exception(trainer, pl_module, exception)#

Called when any trainer execution is interrupted by an exception.

on_fit_end(trainer, pl_module)#

Called when fit ends.

on_fit_start(trainer, pl_module)#

Called when fit begins.

on_load_checkpoint(trainer, pl_module, checkpoint)#

Called when loading a model checkpoint, use to reload state.

Parameters:

trainer (Trainer) – the current Trainer instance.
pl_module (LightningModule) – the current LightningModule instance.
checkpoint (Dict[str, Any]) – the full checkpoint dictionary that got loaded by the Trainer.

on_predict_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)#

Called when the predict batch ends.

on_predict_batch_start(trainer, pl_module, batch, batch_idx, dataloader_idx=0)#

Called when the predict batch begins.

on_predict_end(trainer, pl_module)#

Called when predict ends.

on_predict_epoch_end(trainer, pl_module)#

Called when the predict epoch ends.

on_predict_epoch_start(trainer, pl_module)#

Called when the predict epoch begins.

on_predict_start(trainer, pl_module)#

Called when the predict begins.

on_sanity_check_end(trainer, pl_module)#

Called when the validation sanity check ends.

on_sanity_check_start(trainer, pl_module)#

Called when the validation sanity check starts.

on_save_checkpoint(trainer, pl_module, checkpoint)#

Called when saving a checkpoint to give you a chance to store anything else you might want to save.

Parameters:

trainer (Trainer) – the current Trainer instance.
pl_module (LightningModule) – the current LightningModule instance.
checkpoint (Dict[str, Any]) – the checkpoint dictionary that will be saved.

on_test_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)#

Called when the test batch ends.

on_test_batch_start(trainer, pl_module, batch, batch_idx, dataloader_idx=0)#

Called when the test batch begins.

on_test_end(trainer, pl_module)#

Called when the test ends.

on_test_epoch_end(trainer, pl_module)#

Called when the test epoch ends.

on_test_epoch_start(trainer, pl_module)#

Called when the test epoch begins.

on_test_start(trainer, pl_module)#

Called when the test begins.

on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)#

Called when the train batch ends.

Note

The value outputs["loss"] here will be the normalized value w.r.t accumulate_grad_batches of the loss returned from training_step.

on_train_batch_start(trainer, pl_module, batch, batch_idx)#

Called when the train batch begins.

on_train_end(trainer, pl_module)#

Called when the train ends.

on_train_epoch_end(trainer, pl_module)#

Called when the train epoch ends.

To access all batch outputs at the end of the epoch, you can cache step outputs as an attribute of the lightning.pytorch.core.LightningModule and access them in this hook:

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

class MyCallback(L.Callback):
    def on_train_epoch_end(self, trainer, pl_module):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(pl_module.training_step_outputs).mean()
        pl_module.log("training_epoch_mean", epoch_mean)
        # free up the memory
        pl_module.training_step_outputs.clear()

on_train_epoch_start(trainer, pl_module)#

Called when the train epoch begins.

on_train_start(trainer, pl_module)#

Called when the train begins.

on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)#

Called when the validation batch ends.

on_validation_batch_start(trainer, pl_module, batch, batch_idx, dataloader_idx=0)#

Called when the validation batch begins.

on_validation_end(trainer, pl_module)#

Called when the validation loop ends.

on_validation_epoch_end(trainer, pl_module)#

Called when the val epoch ends.

on_validation_epoch_start(trainer, pl_module)#

Called when the val epoch begins.

on_validation_start(trainer, pl_module)#

Called when the validation loop begins.

setup(trainer, pl_module, stage)#

Called when fit, validate, test, predict, or tune begins.

state_dict()#

Called when saving a checkpoint, implement to generate callback’s state_dict.

Returns:: A dictionary containing callback state.
Return type:: Dict[str, Any]

teardown(trainer, pl_module, stage)#

Called when fit, validate, test, predict, or tune ends.

property state_key: str#

Identifier for the state of the callback.

Used to store and retrieve a callback’s state from the checkpoint dictionary by checkpoint["callbacks"][state_key]. Implementations of a callback need to provide a unique state key if 1) the callback has state and 2) it is desired to maintain the state of multiple instances of that callback.

class DictConfig(content, key=None, parent=None, ref_type=typing.Any, key_type=typing.Any, element_type=typing.Any, is_optional=True, flags=None)#

Bases: BaseContainer, MutableMapping[Any, Any]

__init__(content, key=None, parent=None, ref_type=typing.Any, key_type=typing.Any, element_type=typing.Any, is_optional=True, flags=None)#

copy()#

get(key, default_value=None)#

Return the value for key if key is in the dictionary, else default_value (defaulting to None).

items() → a set-like object providing a view on D's items#

items_ex(resolve=True, keys=None)#

keys() → a set-like object providing a view on D's keys#

pop(k[, d]) → v, remove specified key and return the corresponding value.#

If key is not found, d is returned if given, otherwise KeyError is raised.

setdefault(k[, d]) → D.get(k,d), also set D[k]=d if k not in D#

class LightningModule(*args, **kwargs)#

Bases: _DeviceDtypeModuleMixin, HyperparametersMixin, ModelHooks, DataHooks, CheckpointHooks, Module

__init__(*args, **kwargs)#

all_gather(data, group=None, sync_grads=False)#

Gather tensors or collections of tensors from multiple processes.

This method needs to be called on all processes and the tensors need to have the same shape across all processes, otherwise your program will stall forever.

Parameters:

data (Tensor | Dict | List | Tuple) – int, float, tensor of shape (batch, …), or a (possibly nested) collection thereof.
group (Any | None) – the process group to gather results from. Defaults to all processes (world)
sync_grads (bool) – flag that allows users to synchronize gradients for the all_gather operation

Returns:

A tensor of shape (world_size, batch, …), or if the input was a collection the output will also be a collection with tensors of this shape. For the special case where world_size is 1, no additional dimension is added to the tensor(s).

Return type:

Tensor | Dict | List | Tuple

backward(loss, *args, **kwargs)#

Called to perform backward on the loss returned in training_step(). Override this hook with your own implementation if you need to.

Parameters:: loss (Tensor) – The loss tensor returned by training_step(). If gradient accumulation is used, the loss here holds the normalized value (scaled by 1 / accumulation steps).

Example:

def backward(self, loss):
    loss.backward()

clip_gradients(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)#

Handles gradient clipping internally.

Note

Do not override this method. If you want to customize gradient clipping, consider using configure_gradient_clipping() method.
For manual optimization (self.automatic_optimization = False), if you want to use gradient clipping, consider calling self.clip_gradients(opt, gradient_clip_val=0.5, gradient_clip_algorithm="norm") manually in the training step.

Parameters:

optimizer (Optimizer) – Current optimizer being used.
gradient_clip_val (int | float | None) – The value at which to clip gradients.
gradient_clip_algorithm (str | None) – The gradient clipping algorithm to use. Pass gradient_clip_algorithm="value" to clip by value, and gradient_clip_algorithm="norm" to clip by norm.

configure_callbacks()#

Configure model-specific callbacks. When the model gets attached, e.g., when .fit() or .test() gets called, the list or a callback returned here will be merged with the list of callbacks passed to the Trainer’s callbacks argument. If a callback returned here has the same type as one or several callbacks already present in the Trainer’s callbacks list, it will take priority and replace them. In addition, Lightning will make sure ModelCheckpoint callbacks run last.

Returns:: A callback or a list of callbacks which will extend the list of callbacks in the Trainer.
Return type:: Sequence[Callback] | Callback

Example:

def configure_callbacks(self):
    early_stop = EarlyStopping(monitor="val_acc", mode="max")
    checkpoint = ModelCheckpoint(monitor="val_loss")
    return [early_stop, checkpoint]

configure_gradient_clipping(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)#

Perform gradient clipping for the optimizer parameters. Called before optimizer_step().

Parameters:

optimizer (Optimizer) – Current optimizer being used.
gradient_clip_val (int | float | None) – The value at which to clip gradients. By default, value passed in Trainer will be available here.
gradient_clip_algorithm (str | None) – The gradient clipping algorithm to use. By default, value passed in Trainer will be available here.

Example:

def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm):
    # Implement your own custom logic to clip gradients
    # You can call `self.clip_gradients` with your settings:
    self.clip_gradients(
        optimizer,
        gradient_clip_val=gradient_clip_val,
        gradient_clip_algorithm=gradient_clip_algorithm
    )

configure_optimizers()#

Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you’d need one. But in the case of GANs or similar you might have multiple. Optimization with multiple optimizers only works in the manual optimization mode.

Returns:

Any of these 6 options.

Single optimizer.
List or Tuple of optimizers.
Two lists - The first list has multiple optimizers, and the second has multiple LR schedulers (or multiple lr_scheduler_config).
Dictionary, with an "optimizer" key, and (optionally) a "lr_scheduler" key whose value is a single LR scheduler or lr_scheduler_config.
None - Fit will run without any optimizer.

Return type:

The lr_scheduler_config is a dictionary which contains the scheduler and its associated configuration. The default configuration is shown below.

lr_scheduler_config = {
    # REQUIRED: The scheduler instance
    "scheduler": lr_scheduler,
    # The unit of the scheduler's step size, could also be 'step'.
    # 'epoch' updates the scheduler on epoch end whereas 'step'
    # updates it after a optimizer update.
    "interval": "epoch",
    # How many epochs/steps should pass between calls to
    # `scheduler.step()`. 1 corresponds to updating the learning
    # rate after every epoch/step.
    "frequency": 1,
    # Metric to to monitor for schedulers like `ReduceLROnPlateau`
    "monitor": "val_loss",
    # If set to `True`, will enforce that the value specified 'monitor'
    # is available when the scheduler is updated, thus stopping
    # training if not found. If set to `False`, it will only produce a warning
    "strict": True,
    # If using the `LearningRateMonitor` callback to monitor the
    # learning rate progress, this keyword can be used to specify
    # a custom logged name
    "name": None,
}

When there are schedulers in which the .step() method is conditioned on a value, such as the torch.optim.lr_scheduler.ReduceLROnPlateau scheduler, Lightning requires that the lr_scheduler_config contains the keyword "monitor" set to the metric name that the scheduler should be conditioned on.

# The ReduceLROnPlateau scheduler requires a monitor
def configure_optimizers(self):
    optimizer = Adam(...)
    return {
        "optimizer": optimizer,
        "lr_scheduler": {
            "scheduler": ReduceLROnPlateau(optimizer, ...),
            "monitor": "metric_to_track",
            "frequency": "indicates how often the metric is updated",
            # If "monitor" references validation metrics, then "frequency" should be set to a
            # multiple of "trainer.check_val_every_n_epoch".
        },
    }

# In the case of two optimizers, only one using the ReduceLROnPlateau scheduler
def configure_optimizers(self):
    optimizer1 = Adam(...)
    optimizer2 = SGD(...)
    scheduler1 = ReduceLROnPlateau(optimizer1, ...)
    scheduler2 = LambdaLR(optimizer2, ...)
    return (
        {
            "optimizer": optimizer1,
            "lr_scheduler": {
                "scheduler": scheduler1,
                "monitor": "metric_to_track",
            },
        },
        {"optimizer": optimizer2, "lr_scheduler": scheduler2},
    )

Metrics can be made available to monitor by simply logging it using self.log('metric_to_track', metric_val) in your LightningModule.

Note

Some things to know:

Lightning calls .backward() and .step() automatically in case of automatic optimization.
If a learning rate scheduler is specified in configure_optimizers() with key "interval" (default “epoch”) in the scheduler configuration, Lightning will call the scheduler’s .step() method automatically in case of automatic optimization.
If you use 16-bit precision (precision=16), Lightning will automatically handle the optimizer.
If you use torch.optim.LBFGS, Lightning handles the closure function automatically for you.
If you use multiple optimizers, you will have to switch to ‘manual optimization’ mode and step them yourself.
If you need to control how often the optimizer steps, override the optimizer_step() hook.

forward(*args, **kwargs)#

Same as torch.nn.Module.forward().

Parameters:

*args (Any) – Whatever you decide to pass into the forward method.
**kwargs (Any) – Keyword arguments are also possible.

Returns:

Your model’s output

Return type:

Any

freeze()#

Freeze all params for inference.

Example:

model = MyLightningModule(...)
model.freeze()

load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs)#

Primary way of loading a model from a checkpoint. When Lightning saves a checkpoint it stores the arguments passed to __init__ in the checkpoint under "hyper_parameters".

Any arguments specified through **kwargs will override args stored in "hyper_parameters".

Parameters:

checkpoint_path (str | Path | IO) – Path to checkpoint. This can also be a URL, or file-like object
map_location (device | str | int | Callable[[UntypedStorage, str], UntypedStorage | None] | Dict[device | str | int, device | str | int] | None) – If your checkpoint saved a GPU model and you now load on CPUs or a different number of GPUs, use this to map to the new setup. The behaviour is the same as in torch.load().
hparams_file (str | Path | None) –
Optional path to a .yaml or .csv file with hierarchical structure as in this example:
```
drop_prob: 0.2
dataloader:
    batch_size: 32
```
You most likely won’t need this since Lightning will always save the hyperparameters to the checkpoint. However, if your checkpoint weights don’t have the hyperparameters saved, use this method to pass in a .yaml file with the hparams you’d like to use. These will be converted into a dict and passed into your LightningModule for use.

If your model’s hparams argument is Namespace and .yaml file has hierarchical structure, you need to refactor your model to treat hparams as dict.
strict (bool | None) – Whether to strictly enforce that the keys in checkpoint_path match the keys returned by this module’s state dict. Defaults to True unless LightningModule.strict_loading is set, in which case it defaults to the value of LightningModule.strict_loading.
**kwargs (Any) – Any extra keyword args needed to init the model. Can also be used to override saved hyperparameter values.

Returns:

LightningModule instance with loaded weights and hyperparameters (if available).

Return type:

Self

Note

load_from_checkpoint is a class method. You should use your LightningModule class to call it instead of the LightningModule instance, or a TypeError will be raised.

Note

To ensure all layers can be loaded from the checkpoint, this function will call configure_model() directly after instantiating the model if this hook is overridden in your LightningModule. However, note that load_from_checkpoint does not support loading sharded checkpoints, and you may run out of memory if the model is too large. In this case, consider loading through the Trainer via .fit(ckpt_path=...).

Example:

# load weights without mapping ...
model = MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')

# or load weights mapping all weights from GPU 1 to GPU 0 ...
map_location = {'cuda:1':'cuda:0'}
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    map_location=map_location
)

# or load weights and hyperparameters from separate files.
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    hparams_file='/path/to/hparams_file.yaml'
)

# override some of the params with new values
model = MyLightningModule.load_from_checkpoint(
    PATH,
    num_layers=128,
    pretrained_ckpt_path=NEW_PATH,
)

# predict
pretrained_model.eval()
pretrained_model.freeze()
y_hat = pretrained_model(x)

log(name, value, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, metric_attribute=None, rank_zero_only=False)#

Log a key, value pair.

Example:

self.log('train_loss', loss)

The default behavior per hook is documented here: extensions/logging:Automatic Logging.

Parameters:

name (str) – key to log. Must be identical across all processes if using DDP or any other distributed strategy.
value (Metric | Tensor | int | float) – value to log. Can be a float, Tensor, or a Metric.
prog_bar (bool) – if True logs to the progress bar.
logger (bool | None) – if True logs to the logger.
on_step (bool | None) – if True logs at this step. The default value is determined by the hook. See extensions/logging:Automatic Logging for details.
on_epoch (bool | None) – if True logs epoch accumulated metrics. The default value is determined by the hook. See extensions/logging:Automatic Logging for details.
reduce_fx (str | Callable) – reduction function over step values for end of epoch. torch.mean() by default.
enable_graph (bool) – if True, will not auto detach the graph.
sync_dist (bool) – if True, reduces the metric across devices. Use with care as this may lead to a significant communication overhead.
sync_dist_group (Any | None) – the DDP group to sync across.
add_dataloader_idx (bool) – if True, appends the index of the current dataloader to the name (when using multiple dataloaders). If False, user needs to give unique names for each dataloader to not mix the values.
batch_size (int | None) – Current batch_size. This will be directly inferred from the loaded batch, but for some data structures you might need to explicitly provide it.
metric_attribute (str | None) – To restore the metric state, Lightning requires the reference of the torchmetrics.Metric in your model. This is found automatically if it is a model attribute.
rank_zero_only (bool) – Tells Lightning if you are calling self.log from every process (default) or only from rank 0. If True, you won’t be able to use this metric as a monitor in callbacks (e.g., early stopping). Warning: Improper use can lead to deadlocks! See Advanced Logging for more details.

log_dict(dictionary, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, rank_zero_only=False)#

Log a dictionary of values at once.

Example:

values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n}
self.log_dict(values)

Parameters:

dictionary (Mapping[str, Metric | Tensor | int | float] | MetricCollection) – key value pairs. Keys must be identical across all processes if using DDP or any other distributed strategy. The values can be a float, Tensor, Metric, or MetricCollection.
prog_bar (bool) – if True logs to the progress base.
logger (bool | None) – if True logs to the logger.
on_step (bool | None) – if True logs at this step. None auto-logs for training_step but not validation/test_step. The default value is determined by the hook. See extensions/logging:Automatic Logging for details.
on_epoch (bool | None) – if True logs epoch accumulated metrics. None auto-logs for val/test step but not training_step. The default value is determined by the hook. See extensions/logging:Automatic Logging for details.
reduce_fx (str | Callable) – reduction function over step values for end of epoch. torch.mean() by default.
enable_graph (bool) – if True, will not auto-detach the graph
sync_dist (bool) – if True, reduces the metric across GPUs/TPUs. Use with care as this may lead to a significant communication overhead.
sync_dist_group (Any | None) – the ddp group to sync across.
add_dataloader_idx (bool) – if True, appends the index of the current dataloader to the name (when using multiple). If False, user needs to give unique names for each dataloader to not mix values.
batch_size (int | None) – Current batch size. This will be directly inferred from the loaded batch, but some data structures might need to explicitly provide it.
rank_zero_only (bool) – Tells Lightning if you are calling self.log from every process (default) or only from rank 0. If True, you won’t be able to use this metric as a monitor in callbacks (e.g., early stopping). Warning: Improper use can lead to deadlocks! See Advanced Logging for more details.

lr_scheduler_step(scheduler, metric)#

Override this method to adjust the default way the Trainer calls each scheduler. By default, Lightning calls step() and as shown in the example for each scheduler based on its interval.

Parameters:

scheduler (LRScheduler | ReduceLROnPlateau) – Learning rate scheduler.
metric (Any | None) – Value of the monitor used for schedulers like ReduceLROnPlateau.

Examples:

# DEFAULT
def lr_scheduler_step(self, scheduler, metric):
    if metric is None:
        scheduler.step()
    else:
        scheduler.step(metric)

# Alternative way to update schedulers if it requires an epoch value
def lr_scheduler_step(self, scheduler, metric):
    scheduler.step(epoch=self.current_epoch)

lr_schedulers()#

Returns the learning rate scheduler(s) that are being used during training. Useful for manual optimization.

Returns:: A single scheduler, or a list of schedulers in case multiple ones are present, or None if no schedulers were returned in configure_optimizers().
Return type:: None | List[LRScheduler | ReduceLROnPlateau] | LRScheduler | ReduceLROnPlateau

manual_backward(loss, *args, **kwargs)#

Call this directly from your training_step() when doing optimizations manually. By using this, Lightning can ensure that all the proper scaling gets applied when using mixed precision.

See manual optimization for more examples.

Example:

def training_step(...):
    opt = self.optimizers()
    loss = ...
    opt.zero_grad()
    # automatically applies scaling, etc...
    self.manual_backward(loss)
    opt.step()

Parameters:

loss (Tensor) – The tensor on which to compute gradients. Must have a graph attached.
*args (Any) – Additional positional arguments to be forwarded to backward()
**kwargs (Any) – Additional keyword arguments to be forwarded to backward()

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)#

Override this method to adjust the default way the Trainer calls the optimizer.

By default, Lightning calls step() and zero_grad() as shown in the example. This method (and zero_grad()) won’t be called during the accumulation phase when Trainer(accumulate_grad_batches != 1). Overriding this hook has no benefit with manual optimization.

Parameters:

epoch (int) – Current epoch
batch_idx (int) – Index of current batch
optimizer (Optimizer | LightningOptimizer) – A PyTorch optimizer
optimizer_closure (Callable[[], Any] | None) – The optimizer closure. This closure must be executed as it includes the calls to training_step(), optimizer.zero_grad(), and backward().

Examples:

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    # Add your custom logic to run directly before `optimizer.step()`

    optimizer.step(closure=optimizer_closure)

    # Add your custom logic to run directly after `optimizer.step()`

optimizer_zero_grad(epoch, batch_idx, optimizer)#

Override this method to change the default behaviour of optimizer.zero_grad().

Parameters:

epoch (int) – Current epoch
batch_idx (int) – Index of current batch
optimizer (Optimizer) – A PyTorch optimizer

Examples:

# DEFAULT
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad()

# Set gradients to `None` instead of zero to improve performance (not required on `torch>=2.0.0`).
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad(set_to_none=True)

See torch.optim.Optimizer.zero_grad() for the explanation of the above example.

optimizers(use_pl_optimizer: Literal[True] = True) → LightningOptimizer | List[LightningOptimizer]#

optimizers(use_pl_optimizer: Literal[False]) → Optimizer | List[Optimizer]

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

Parameters:: use_pl_optimizer – If True, will wrap the optimizer(s) in a LightningOptimizer for automatic handling of precision, profiling, and counting of step calls for proper logging and checkpointing. It specifically wraps the step method and custom optimizers that don’t have this method are not supported.
Returns:: A single optimizer, or a list of optimizers in case multiple ones are present.

predict_step(*args, **kwargs)#

Step function called during predict(). By default, it calls forward(). Override to add any processing logic.

The predict_step() is used to scale inference on multi-devices.

To prevent an OOM error, it is possible to use BasePredictionWriter callback to write the predictions to disk or database after each batch or on epoch end.

The BasePredictionWriter should be used while using a spawn based accelerator. This happens for Trainer(strategy="ddp_spawn") or training on 8 TPU cores with Trainer(accelerator="tpu", devices=8) as predictions won’t be returned.

Parameters:

batch – The output of your data iterable, normally a DataLoader.
batch_idx – The index of this batch.
dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

Predicted output (optional).

Return type:

Any

Example

class MyModel(LightningModule):

    def predict_step(self, batch, batch_idx, dataloader_idx=0):
        return self(batch)

dm = ...
model = MyModel()
trainer = Trainer(accelerator="gpu", devices=2)
predictions = trainer.predict(model, dm)

print(*args, **kwargs)#

Prints only from process 0. Use this in any distributed mode to log only once.

Parameters:

*args (Any) – The thing to print. The same as for Python’s built-in print function.
**kwargs (Any) – The same as for Python’s built-in print function.

Example:

def forward(self, x):
    self.print(x, 'in forward')

test_step(*args, **kwargs)#

Operates on a single batch of data from the test set. In this step you’d normally generate examples or calculate anything of interest such as accuracy.

Parameters:

batch – The output of your data iterable, normally a DataLoader.
batch_idx – The index of this batch.
dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

Tensor - The loss tensor
dict - A dictionary. Can include any keys, but must include the key 'loss'.
None - Skip to the next batch.

Return type:

Tensor | Mapping[str, Any] | None

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...

# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

If you pass in multiple test dataloaders, test_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

Note

If you don’t need to test you don’t need to implement this method.

Note

When the test_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of the test epoch, the model goes back to training mode and gradients are enabled.

to_onnx(file_path, input_sample=None, **kwargs)#

Saves the model in ONNX format.

Parameters:

file_path (str | Path) – The path of the file the onnx model should be saved to.
input_sample (Any | None) – An input for tracing. Default: None (Use self.example_input_array)
**kwargs (Any) – Will be passed to torch.onnx.export function.

Example:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)

model = SimpleModel()
input_sample = torch.randn(1, 64)
model.to_onnx("export.onnx", input_sample, export_params=True)

to_torchscript(file_path=None, method='script', example_inputs=None, **kwargs)#

By default compiles the whole model to a ScriptModule. If you want to use tracing, please provided the argument method='trace' and make sure that either the example_inputs argument is provided, or the model has example_input_array set. If you would like to customize the modules that are scripted you should override this method. In case you want to return multiple modules, we recommend using a dictionary.

Parameters:

file_path (str | Path | None) – Path where to save the torchscript. Default: None (no file saved).
method (str | None) – Whether to use TorchScript’s script or trace method. Default: ‘script’
example_inputs (Any | None) – An input to be used to do tracing when method is set to ‘trace’. Default: None (uses example_input_array)
**kwargs (Any) – Additional arguments that will be passed to the torch.jit.script() or torch.jit.trace() function.

Note

Requires the implementation of the forward() method.
The exported script will be set to evaluation mode.
It is recommended that you install the latest supported version of PyTorch to use this feature without limitations. See also the torch.jit documentation for supported features.

Example:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

model = SimpleModel()
model.to_torchscript(file_path="model.pt")

torch.jit.save(model.to_torchscript(
    file_path="model_trace.pt", method='trace', example_inputs=torch.randn(1, 64))
)

Returns:: This LightningModule as a torchscript, regardless of whether file_path is defined or not.
Return type:: ScriptModule | Dict[str, ScriptModule]

toggle_optimizer(optimizer)#

Makes sure only the gradients of the current optimizer’s parameters are calculated in the training step to prevent dangling gradients in multiple-optimizer setup.

It works with untoggle_optimizer() to make sure param_requires_grad_state is properly reset.

Parameters:: optimizer (Optimizer | LightningOptimizer) – The optimizer to toggle.

training_step(*args, **kwargs)#

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Parameters:

batch – The output of your data iterable, normally a DataLoader.
batch_idx – The index of this batch.
dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

Tensor - The loss tensor
dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.
None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

Return type:

Tensor | Mapping[str, Any] | None

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False

# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note

When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

unfreeze()#

Unfreeze all parameters for training.

model = MyLightningModule(...)
model.unfreeze()

untoggle_optimizer(optimizer)#

Resets the state of required gradients that were toggled with toggle_optimizer().

Parameters:: optimizer (Optimizer | LightningOptimizer) – The optimizer to untoggle.

validation_step(*args, **kwargs)#

Operates on a single batch of data from the validation set. In this step you’d might generate examples or calculate anything of interest like accuracy.

Parameters:

batch – The output of your data iterable, normally a DataLoader.
batch_idx – The index of this batch.
dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

Tensor - The loss tensor
dict - A dictionary. Can include any keys, but must include the key 'loss'.
None - Skip to the next batch.

Return type:

Tensor | Mapping[str, Any] | None

# if you have one val dataloader:
def validation_step(self, batch, batch_idx): ...

# if you have multiple val dataloaders:
def validation_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single validation dataset
def validation_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'val_loss': loss, 'val_acc': val_acc})

If you pass in multiple val dataloaders, validation_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple validation dataloaders
def validation_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

Note

If you don’t need to validate you don’t need to implement this method.

Note

When the validation_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, the model goes back to training mode and gradients are enabled.

CHECKPOINT_HYPER_PARAMS_KEY = 'hyper_parameters'#

CHECKPOINT_HYPER_PARAMS_NAME = 'hparams_name'#

CHECKPOINT_HYPER_PARAMS_TYPE = 'hparams_type'#

allow_zero_length_dataloader_with_multiple_devices: bool#

property automatic_optimization: bool#: If set to False you are responsible for calling .backward(), .step(), .zero_grad().

property current_epoch: int#: The current epoch in the Trainer, or 0 if not attached.

property device_mesh: DeviceMesh | None#: Strategies like ModelParallelStrategy will create a device mesh that can be accessed in the configure_model() hook to parallelize the LightningModule.

property example_input_array: Tensor | Tuple | Dict | None#

The example input array is a specification of what the module can consume in the forward() method. The return type is interpreted as follows:

Single tensor: It is assumed the model takes a single argument, i.e., model.forward(model.example_input_array)
Tuple: The input array should be interpreted as a sequence of positional arguments, i.e., model.forward(*model.example_input_array)
Dict: The input array represents named keyword arguments, i.e., model.forward(**model.example_input_array)

property fabric: Fabric | None#: !! processed by numpydoc !!

property global_rank: int#: The index of the current process across all nodes and devices.

property global_step: int#

Total training batches seen across all epochs.

If no Trainer is attached, this propery is 0.

property local_rank: int#: The index of the current process within a single node.

property logger: Logger | Logger | None#: Reference to the logger object in the Trainer.

property loggers: List[Logger] | List[Logger]#: Reference to the list of loggers in the Trainer.

property on_gpu: bool#

Returns True if this model is currently located on a GPU.

Useful to set flags around the LightningModule for different CPU vs GPU behavior.

prepare_data_per_node: bool#

property strict_loading: bool#: Determines how Lightning loads this model using .load_state_dict(…, strict=model.strict_loading).

property trainer: Trainer#: !! processed by numpydoc !!

training: bool#

class Logger#

Bases: Logger, ABC

Base class for experiment loggers.

after_save_checkpoint(checkpoint_callback)#

Called after model checkpoint callback saves a new checkpoint.

Parameters:: checkpoint_callback (ModelCheckpoint) – the model checkpoint callback instance

property save_dir: str | None#: Return the root directory where experiment logs get saved, or None if the logger does not save data locally.

class OmegaConf#

Bases: object

OmegaConf primary class

__init__()#

static clear_cache(conf)#

classmethod clear_resolver(name)#

Clear(remove) any resolver only if it exists.

Returns a bool: True if resolver is removed and False if not removed.

Parameters:: name (str) – Name of the resolver.
Returns:: A bool (True if resolver is removed, False if not found before removing).
Return type:: bool

static clear_resolvers()#

Clear(remove) all OmegaConf resolvers, then re-register OmegaConf’s default resolvers.

static copy_cache(from_config, to_config)#

static create(obj: str, parent: BaseContainer | None = None, flags: Dict[str, bool] | None = None) → DictConfig | ListConfig#
static create(obj: List[Any] | Tuple[Any, ...], parent: BaseContainer | None = None, flags: Dict[str, bool] | None = None) → ListConfig
static create(obj: DictConfig, parent: BaseContainer | None = None, flags: Dict[str, bool] | None = None) → DictConfig
static create(obj: ListConfig, parent: BaseContainer | None = None, flags: Dict[str, bool] | None = None) → ListConfig
static create(obj: Dict[Any, Any] | None = None, parent: BaseContainer | None = None, flags: Dict[str, bool] | None = None) → DictConfig

static from_cli(args_list=None)#

static from_dotlist(dotlist)#

Creates config from the content sys.argv or from the specified args list of not None

Parameters:: dotlist (List[str]) – A list of dotlist-style strings, e.g. ["foo.bar=1", "baz=qux"].
Returns:: A DictConfig object created from the dotlist.
Return type:: DictConfig

static get_cache(conf)#

static get_type(obj, key=None)#

classmethod has_resolver(name)#

static is_config(obj)#

static is_dict(obj)#

static is_interpolation(node, key=None)#

static is_list(obj)#

static is_missing(cfg, key)#

static is_readonly(conf)#

static is_struct(conf)#

static legacy_register_resolver(name, resolver)#

static load(file_)#

static masked_copy(conf, keys)#

Create a masked copy of of this config that contains a subset of the keys

Parameters:

conf (DictConfig) – DictConfig object
keys (str | List[str]) – keys to preserve in the copy

Returns:

The masked DictConfig object.

Return type:

DictConfig

static merge(*configs)#

Merge a list of previously created configs into a single one

Parameters:: configs (DictConfig | ListConfig | Dict[str | bytes | int | Enum | float | bool, Any] | List[Any] | Tuple[Any, ...] | Any) – Input configs
Returns:: the merged config object.
Return type:: ListConfig | DictConfig

static missing_keys(cfg)#

Returns a set of missing keys in a dotlist style.

Parameters:: cfg (Any) – An OmegaConf.Container, or a convertible object via OmegaConf.create (dict, list, …).
Returns:: set of strings of the missing keys.
Raises:: ValueError – On input not representing a config.
Return type:: Set[str]

static register_new_resolver(name, resolver, *, replace=False, use_cache=False)#

Parameters:

name (str) – Name of the resolver.
resolver (Callable[[...], Any]) – Callable whose arguments are provided in the interpolation, e.g., with ${foo:x,0,${y.z}} these arguments are respectively “x” (str), 0 (int) and the value of y.z.
replace (bool) – If set to False (default), then a ValueError is raised if an existing resolver has already been registered with the same name. If set to True, then the new resolver replaces the previous one. NOTE: The cache on existing config objects is not affected, use OmegaConf.clear_cache(cfg) to clear it.
use_cache (bool) – Whether the resolver’s outputs should be cached. The cache is based only on the string literals representing the resolver arguments, e.g., ${foo:${bar}} will always return the same value regardless of the value of bar if the cache is enabled for foo.

static register_resolver(name, resolver)#

static resolve(cfg)#

Resolves all interpolations in the given config object in-place.

Parameters:: cfg (Container) – An OmegaConf container (DictConfig, ListConfig) Raises a ValueError if the input object is not an OmegaConf container.

static save(config, f, resolve=False)#

Save as configuration object to a file

Parameters:

config (Any) – omegaconf.Config object (DictConfig or ListConfig).
f (str | Path | IO[Any]) – filename or file object
resolve (bool) – True to save a resolved config (defaults to False)

static select(cfg, key, *, default=_DEFAULT_MARKER_, throw_on_resolution_failure=True, throw_on_missing=False)#

Parameters:

cfg (Container) – Config node to select from
key (str) – Key to select
default (Any) – Default value to return if key is not found
throw_on_resolution_failure (bool) – Raise an exception if an interpolation resolution error occurs, otherwise return None
throw_on_missing (bool) – Raise an exception if an attempt to select a missing key (with the value ‘???’) is made, otherwise return None

Returns:

selected value or None if not found.

Return type:

Any

static set_cache(conf, cache)#

static set_readonly(conf, value)#

static set_struct(conf, value)#

static structured(obj, parent=None, flags=None)#

static to_container(cfg, *, resolve=False, throw_on_missing=False, enum_to_str=False, structured_config_mode=SCMode.DICT)#

Resursively converts an OmegaConf config to a primitive container (dict or list).

Parameters:

cfg (Any) – the config to convert
resolve (bool) – True to resolve all values
throw_on_missing (bool) – When True, raise MissingMandatoryValue if any missing values are present. When False (the default), replace missing values with the string “???” in the output container.
enum_to_str (bool) – True to convert Enum keys and values to strings
structured_config_mode (SCMode) –
Specify how Structured Configs (DictConfigs backed by a dataclass) are handled.
- By default (structured_config_mode=SCMode.DICT) structured configs are converted to plain dicts.
- If structured_config_mode=SCMode.DICT_CONFIG, structured config nodes will remain as DictConfig.
- If structured_config_mode=SCMode.INSTANTIATE, this function will instantiate structured configs (DictConfigs backed by a dataclass), by creating an instance of the underlying dataclass.
See also OmegaConf.to_object.

Returns:

A dict or a list representing this config as a primitive container.

Return type:

static to_object(cfg)#

Resursively converts an OmegaConf config to a primitive container (dict or list). Any DictConfig objects backed by dataclasses or attrs classes are instantiated as instances of those backing classes.

This is an alias for OmegaConf.to_container(…, resolve=True, throw_on_missing=True,: structured_config_mode=SCMode.INSTANTIATE)

Parameters:: cfg (Any) – the config to convert
Returns:: A dict or a list or dataclass representing this config.
Return type:: Dict[str | bytes | int | Enum | float | bool, Any] | List[Any] | None | str | Any

static to_yaml(cfg, *, resolve=False, sort_keys=False)#

returns a yaml dump of this config object.

Parameters:

cfg (Any) – Config object, Structured Config type or instance
resolve (bool) – if True, will return a string with the interpolations resolved, otherwise interpolations are preserved
sort_keys (bool) – If True, will print dict keys in sorted order. default False.

Returns:

A string containing the yaml representation.

Return type:

str

static unsafe_merge(*configs)#

Merge a list of previously created configs into a single one This is much faster than OmegaConf.merge() as the input configs are not copied. However, the input configs must not be used after this operation as will become inconsistent.

Parameters:: configs (DictConfig | ListConfig | Dict[str | bytes | int | Enum | float | bool, Any] | List[Any] | Tuple[Any, ...] | Any) – Input configs
Returns:: the merged config object.
Return type:: ListConfig | DictConfig

static update(cfg, key, value=None, *, merge=True, force_add=False)#

Updates a dot separated key sequence to a value

Parameters:

cfg (Container) – input config to update
key (str) – key to update (can be a dot separated path)
value (Any) – value to set, if value if a list or a dict it will be merged or set depending on merge_config_values
merge (bool) – If value is a dict or a list, True (default) to merge into the destination, False to replace the destination.
force_add (bool) – insert the entire path regardless of Struct flag or Structured Config nodes.

class PreProcessor(dataset, data_dir, transforms_config=None, **kwargs)#

Bases: InMemoryDataset

Preprocessor for datasets.

Parameters:

datasetlist: List of data objects.
data_dirstr: Path to the directory containing the data.
transforms_configDictConfig, optional: Configuration parameters for the transforms (default: None).
**kwargsoptional: Optional additional arguments.

__init__(dataset, data_dir, transforms_config=None, **kwargs)#

instantiate_pre_transform(data_dir, transforms_config)#

Instantiate the pre-transforms.

Parameters:

data_dirstr: Path to the directory containing the data.
transforms_configDictConfig: Configuration parameters for the transforms.

Returns:

torch_geometric.transforms.Compose: Pre-transform object.

load(path)#

Load the dataset from the file path path.

Parameters:

pathstr: The path to the processed data.

load_dataset_splits(split_params)#

Load the dataset splits.

Parameters:

split_paramsdict: Parameters for loading the dataset splits.

Returns:

tuple: A tuple containing the train, validation, and test datasets.

process()#

Method that processes the data.

save_transform_parameters()#

Save the transform parameters.

set_processed_data_dir(pre_transforms_dict, data_dir, transforms_config)#

Set the processed data directory.

Parameters:

pre_transforms_dictdict: Dictionary containing the pre-transforms.
data_dirstr: Path to the directory containing the data.
transforms_configDictConfig: Configuration parameters for the transforms.

property processed_dir: str#

Return the path to the processed directory.

Returns:

str: Path to the processed directory.

property processed_file_names: str#

Return the name of the processed file.

Returns:

str: Name of the processed file.

class RankedLogger(name='topobench.utils.pylogger', rank_zero_only=False, extra=None)#

Bases: LoggerAdapter

Initialize a multi-GPU-friendly python command line logger.

The logger logs on all processes with their rank prefixed in the log message.

Parameters:

namestr, optional: The name of the logger, by default __name__.
rank_zero_onlybool, optional: Whether to force all logs to only occur on the rank zero process (default: False).
extraMapping[str, object], optional: A dict-like object which provides contextual information. See logging.LoggerAdapter for more information (default: None).

__init__(name='topobench.utils.pylogger', rank_zero_only=False, extra=None)#

Initialize the adapter with a logger and a dict-like object which provides contextual information. This constructor signature allows easy stacking of LoggerAdapters, if so desired.

You can effectively pass keyword arguments as shown in the following example:

adapter = LoggerAdapter(someLogger, dict(p1=v1, p2=”v2”))

log(level, msg, rank=None, *args, **kwargs)#

Delegate a log call to the underlying logger.

The function first prefixes the message with the rank of the process it’s being logged from and then logs the message. If ‘rank’ is provided, then the log will only occur on that rank/process.

Parameters:

levelint: The level to log at. Look at logging.__init__.py for more information.
msgstr: The message to log.
rankint, optional: The rank to log at (default: None).
*argsAny: Additional args to pass to the underlying logging function.
**kwargsAny: Any additional keyword args to pass to the underlying logging function.

class TBDataloader(dataset_train, dataset_val=None, dataset_test=None, batch_size=1, num_workers=0, pin_memory=False, **kwargs)#

Bases: LightningDataModule

This class takes care of returning the dataloaders for the training, validation, and test datasets.

It also handles the collate function. The class is designed to work with the torch dataloaders.

Parameters:

dataset_trainDataloadDataset: The training dataset.
dataset_valDataloadDataset, optional: The validation dataset (default: None).
dataset_testDataloadDataset, optional: The test dataset (default: None).
batch_sizeint, optional: The batch size for the dataloader (default: 1).
num_workersint, optional: The number of worker processes to use for data loading (default: 0).
pin_memorybool, optional: If True, the data loader will copy tensors into pinned memory before returning them (default: False).
**kwargsoptional: Additional arguments.

References

Read the docs:: https://lightning.ai/docs/pytorch/latest/data/datamodule.html

__init__(dataset_train, dataset_val=None, dataset_test=None, batch_size=1, num_workers=0, pin_memory=False, **kwargs)#

prepare_data_per_node#: If True, each LOCAL_RANK=0 will call prepare data. Otherwise only NODE_RANK=0, LOCAL_RANK=0 will prepare data.

allow_zero_length_dataloader_with_multiple_devices#: If True, dataloader with zero length within local rank is allowed. Default value is False.

state_dict()#

Called when saving a checkpoint. Implement to generate and save the datamodule state.

Returns:

dict: A dictionary containing the datamodule state that you want to save.

teardown(stage=None)#

Lightning hook for cleaning up after trainer.fit(), trainer.validate(), trainer.test(), and trainer.predict().

Parameters:

stagestr, optional: The stage being torn down. Either “fit”, “validate”, “test”, or “predict” (default: None).

test_dataloader()#

Create and return the test dataloader.

Returns:

torch.utils.data.DataLoader: The test dataloader.

train_dataloader()#

Create and return the train dataloader.

Returns:

torch.utils.data.DataLoader: The train dataloader.

val_dataloader()#

Create and return the validation dataloader.

Returns:

torch.utils.data.DataLoader: The validation dataloader.

class Trainer(*, accelerator='auto', strategy='auto', devices='auto', num_nodes=1, precision=None, logger=None, callbacks=None, fast_dev_run=False, max_epochs=None, min_epochs=None, max_steps=-1, min_steps=None, max_time=None, limit_train_batches=None, limit_val_batches=None, limit_test_batches=None, limit_predict_batches=None, overfit_batches=0.0, val_check_interval=None, check_val_every_n_epoch=1, num_sanity_val_steps=None, log_every_n_steps=None, enable_checkpointing=None, enable_progress_bar=None, enable_model_summary=None, accumulate_grad_batches=1, gradient_clip_val=None, gradient_clip_algorithm=None, deterministic=None, benchmark=None, inference_mode=True, use_distributed_sampler=True, profiler=None, detect_anomaly=False, barebones=False, plugins=None, sync_batchnorm=False, reload_dataloaders_every_n_epochs=0, default_root_dir=None)#

Bases: object

__init__(*, accelerator='auto', strategy='auto', devices='auto', num_nodes=1, precision=None, logger=None, callbacks=None, fast_dev_run=False, max_epochs=None, min_epochs=None, max_steps=-1, min_steps=None, max_time=None, limit_train_batches=None, limit_val_batches=None, limit_test_batches=None, limit_predict_batches=None, overfit_batches=0.0, val_check_interval=None, check_val_every_n_epoch=1, num_sanity_val_steps=None, log_every_n_steps=None, enable_checkpointing=None, enable_progress_bar=None, enable_model_summary=None, accumulate_grad_batches=1, gradient_clip_val=None, gradient_clip_algorithm=None, deterministic=None, benchmark=None, inference_mode=True, use_distributed_sampler=True, profiler=None, detect_anomaly=False, barebones=False, plugins=None, sync_batchnorm=False, reload_dataloaders_every_n_epochs=0, default_root_dir=None)#

Customize every aspect of training via flags.

Parameters:

accelerator (str | Accelerator) – Supports passing different accelerator types (“cpu”, “gpu”, “tpu”, “hpu”, “mps”, “auto”) as well as custom accelerator instances.
strategy (str | Strategy) – Supports different training strategies with aliases as well custom strategies. Default: "auto".
devices (List[int] | str | int) – The devices to use. Can be set to a positive number (int or str), a sequence of device indices (list or str), the value -1 to indicate all available devices should be used, or "auto" for automatic selection based on the chosen accelerator. Default: "auto".
num_nodes (int) – Number of GPU nodes for distributed training. Default: 1.
precision (Literal[64, 32, 16] | ~typing.Literal['transformer-engine', 'transformer-engine-float16', '16-true', '16-mixed', 'bf16-true', 'bf16-mixed', '32-true', '64-true'] | ~typing.Literal['64', '32', '16', 'bf16'] | None) – Double precision (64, ‘64’ or ‘64-true’), full precision (32, ‘32’ or ‘32-true’), 16bit mixed precision (16, ‘16’, ‘16-mixed’) or bfloat16 mixed precision (‘bf16’, ‘bf16-mixed’). Can be used on CPU, GPU, TPUs, or HPUs. Default: '32-true'.
logger (Logger | Iterable[Logger] | bool | None) – Logger (or iterable collection of loggers) for experiment tracking. A True value uses the default TensorBoardLogger if it is installed, otherwise CSVLogger. False will disable logging. If multiple loggers are provided, local files (checkpoints, profiler traces, etc.) are saved in the log_dir of the first logger. Default: True.
callbacks (List[Callback] | Callback | None) – Add a callback or list of callbacks. Default: None.
fast_dev_run (int | bool) – Runs n if set to n (int) else 1 if set to True batch(es) of train, val and test to find any bugs (ie: a sort of unit test). Default: False.
max_epochs (int | None) – Stop training once this number of epochs is reached. Disabled by default (None). If both max_epochs and max_steps are not specified, defaults to max_epochs = 1000. To enable infinite training, set max_epochs = -1.
min_epochs (int | None) – Force training for at least these many epochs. Disabled by default (None).
max_steps (int) – Stop training after this number of steps. Disabled by default (-1). If max_steps = -1 and max_epochs = None, will default to max_epochs = 1000. To enable infinite training, set max_epochs to -1.
min_steps (int | None) – Force training for at least these number of steps. Disabled by default (None).
max_time (str | timedelta | Dict[str, int] | None) – Stop training after this amount of time has passed. Disabled by default (None). The time duration can be specified in the format DD:HH:MM:SS (days, hours, minutes seconds), as a datetime.timedelta, or a dictionary with keys that will be passed to datetime.timedelta.
limit_train_batches (int | float | None) – How much of training dataset to check (float = fraction, int = num_batches). Default: 1.0.
limit_val_batches (int | float | None) – How much of validation dataset to check (float = fraction, int = num_batches). Default: 1.0.
limit_test_batches (int | float | None) – How much of test dataset to check (float = fraction, int = num_batches). Default: 1.0.
limit_predict_batches (int | float | None) – How much of prediction dataset to check (float = fraction, int = num_batches). Default: 1.0.
overfit_batches (int | float) – Overfit a fraction of training/validation data (float) or a set number of batches (int). Default: 0.0.
val_check_interval (int | float | None) – How often to check the validation set. Pass a float in the range [0.0, 1.0] to check after a fraction of the training epoch. Pass an int to check after a fixed number of training batches. An int value can only be higher than the number of training batches when check_val_every_n_epoch=None, which validates after every N training batches across epochs or during iteration-based training. Default: 1.0.
check_val_every_n_epoch (int | None) – Perform a validation loop after every N training epochs. If None, validation will be done solely based on the number of training batches, requiring val_check_interval to be an integer value. Default: 1.
num_sanity_val_steps (int | None) – Sanity check runs n validation batches before starting the training routine. Set it to -1 to run all batches in all validation dataloaders. Default: 2.
log_every_n_steps (int | None) – How often to log within steps. Default: 50.
enable_checkpointing (bool | None) – If True, enable checkpointing. It will configure a default ModelCheckpoint callback if there is no user-defined ModelCheckpoint in :paramref:`~lightning.pytorch.trainer.trainer.Trainer.callbacks`. Default: True.
enable_progress_bar (bool | None) – Whether to enable to progress bar by default. Default: True.
enable_model_summary (bool | None) – Whether to enable model summarization by default. Default: True.
accumulate_grad_batches (int) – Accumulates gradients over k batches before stepping the optimizer. Default: 1.
gradient_clip_val (int | float | None) – The value at which to clip gradients. Passing gradient_clip_val=None disables gradient clipping. If using Automatic Mixed Precision (AMP), the gradients will be unscaled before. Default: None.
gradient_clip_algorithm (str | None) – The gradient clipping algorithm to use. Pass gradient_clip_algorithm="value" to clip by value, and gradient_clip_algorithm="norm" to clip by norm. By default it will be set to "norm".
deterministic (bool | Literal['warn'] | None) – If True, sets whether PyTorch operations must use deterministic algorithms. Set to "warn" to use deterministic algorithms whenever possible, throwing warnings on operations that don’t support deterministic mode. If not set, defaults to False. Default: None.
benchmark (bool | None) – The value (True or False) to set torch.backends.cudnn.benchmark to. The value for torch.backends.cudnn.benchmark set in the current session will be used (False if not manually set). If :paramref:`~lightning.pytorch.trainer.trainer.Trainer.deterministic` is set to True, this will default to False. Override to manually set a different value. Default: None.
inference_mode (bool) – Whether to use torch.inference_mode() or torch.no_grad() during evaluation (validate/test/predict).
use_distributed_sampler (bool) – Whether to wrap the DataLoader’s sampler with torch.utils.data.DistributedSampler. If not specified this is toggled automatically for strategies that require it. By default, it will add shuffle=True for the train sampler and shuffle=False for validation/test/predict samplers. If you want to disable this logic, you can pass False and add your own distributed sampler in the dataloader hooks. If True and a distributed sampler was already added, Lightning will not replace the existing one. For iterable-style datasets, we don’t do this automatically.
profiler (Profiler | str | None) – To profile individual steps during training and assist in identifying bottlenecks. Default: None.
detect_anomaly (bool) – Enable anomaly detection for the autograd engine. Default: False.
barebones (bool) – Whether to run in “barebones mode”, where all features that may impact raw speed are disabled. This is meant for analyzing the Trainer overhead and is discouraged during regular training runs. The following features are deactivated: :paramref:`~lightning.pytorch.trainer.trainer.Trainer.enable_checkpointing`, :paramref:`~lightning.pytorch.trainer.trainer.Trainer.logger`, :paramref:`~lightning.pytorch.trainer.trainer.Trainer.enable_progress_bar`, :paramref:`~lightning.pytorch.trainer.trainer.Trainer.log_every_n_steps`, :paramref:`~lightning.pytorch.trainer.trainer.Trainer.enable_model_summary`, :paramref:`~lightning.pytorch.trainer.trainer.Trainer.num_sanity_val_steps`, :paramref:`~lightning.pytorch.trainer.trainer.Trainer.fast_dev_run`, :paramref:`~lightning.pytorch.trainer.trainer.Trainer.detect_anomaly`, :paramref:`~lightning.pytorch.trainer.trainer.Trainer.profiler`, log(), log_dict().
plugins (Precision | ClusterEnvironment | CheckpointIO | LayerSync | List[Precision | ClusterEnvironment | CheckpointIO | LayerSync] | None) – Plugins allow modification of core behavior like ddp and amp, and enable custom lightning plugins. Default: None.
sync_batchnorm (bool) – Synchronize batch norm layers between process groups/whole world. Default: False.
reload_dataloaders_every_n_epochs (int) – Set to a positive integer to reload dataloaders every n epochs. Default: 0.
default_root_dir (str | Path | None) – Default path for logs and weights when no logger/ckpt_callback passed. Default: os.getcwd(). Can be remote file paths such as s3://mybucket/path or ‘hdfs://path/’

Raises:

TypeError – If gradient_clip_val is not an int or float.
MisconfigurationException – If gradient_clip_algorithm is invalid.

fit(model, train_dataloaders=None, val_dataloaders=None, datamodule=None, ckpt_path=None)#

Runs the full optimization routine.

Parameters:

model (LightningModule) – Model to fit.
train_dataloaders (Any | LightningDataModule | None) – An iterable or collection of iterables specifying training samples. Alternatively, a LightningDataModule that defines the train_dataloader hook.
val_dataloaders (Any | None) – An iterable or collection of iterables specifying validation samples.
datamodule (LightningDataModule | None) – A LightningDataModule that defines the train_dataloader hook.
ckpt_path (str | Path | None) – Path/URL of the checkpoint from which training is resumed. Could also be one of two special keywords "last" and "hpc". If there is no checkpoint file at the path, an exception is raised.

Raises:

TypeError – If model is not LightningModule for torch version less than 2.0.0 and if model is not LightningModule or torch._dynamo.OptimizedModule for torch versions greater than or equal to 2.0.0 .

For more information about multiple dataloaders, see this section.

init_module(empty_init=None)#

Tensors that you instantiate under this context manager will be created on the device right away and have the right data type depending on the precision setting in the Trainer.

The parameters and tensors get created on the device and with the right data type right away without wasting memory being allocated unnecessarily.

Parameters:: empty_init (bool | None) – Whether to initialize the model with empty weights (uninitialized memory). If None, the strategy will decide. Some strategies may not support all options. Set this to True if you are loading a checkpoint into a large model.

predict(model=None, dataloaders=None, datamodule=None, return_predictions=None, ckpt_path=None)#

Run inference on your data. This will call the model forward function to compute predictions. Useful to perform distributed and batched predictions. Logging is disabled in the predict hooks.

Parameters:

model (LightningModule | None) – The model to predict with.
dataloaders (Any | LightningDataModule | None) – An iterable or collection of iterables specifying predict samples. Alternatively, a LightningDataModule that defines the predict_dataloader hook.
datamodule (LightningDataModule | None) – A LightningDataModule that defines the predict_dataloader hook.
return_predictions (bool | None) – Whether to return predictions. True by default except when an accelerator that spawns processes is used (not supported).
ckpt_path (str | Path | None) – Either "best", "last", "hpc" or path to the checkpoint you wish to predict. If None and the model instance was passed, use the current weights. Otherwise, the best model checkpoint from the previous trainer.fit call will be loaded if a checkpoint callback is configured.

For more information about multiple dataloaders, see this section.

Returns:

Returns a list of dictionaries, one for each provided dataloader containing their respective predictions.

Raises:

TypeError – If no model is passed and there was no LightningModule passed in the previous run. If model passed is not LightningModule or torch._dynamo.OptimizedModule.
MisconfigurationException – If both dataloaders and datamodule are passed. Pass only one of these.
RuntimeError – If a compiled model is passed and the strategy is not supported.

Return type:

List[Any] | List[List[Any]] | None

See Lightning inference section for more.

print(*args, **kwargs)#

Print something only on the first process. If running on multiple machines, it will print from the first process in each machine.

Arguments passed to this method are forwarded to the Python built-in print() function.

save_checkpoint(filepath, weights_only=False, storage_options=None)#

Runs routine to create a checkpoint.

This method needs to be called on all processes in case the selected strategy is handling distributed checkpointing.

Parameters:

filepath (str | Path) – Path where checkpoint is saved.
weights_only (bool) – If True, will only save the model weights.
storage_options (Any | None) – parameter for how to save to storage, passed to CheckpointIO plugin

Raises:

AttributeError – If the model is not attached to the Trainer before calling this method.

test(model=None, dataloaders=None, ckpt_path=None, verbose=True, datamodule=None)#

Perform one evaluation epoch over the test set. It’s separated from fit to make sure you never run on your test set until you want to.

Parameters:

model (LightningModule | None) – The model to test.
dataloaders (Any | LightningDataModule | None) – An iterable or collection of iterables specifying test samples. Alternatively, a LightningDataModule that defines the test_dataloader hook.
ckpt_path (str | Path | None) – Either "best", "last", "hpc" or path to the checkpoint you wish to test. If None and the model instance was passed, use the current weights. Otherwise, the best model checkpoint from the previous trainer.fit call will be loaded if a checkpoint callback is configured.
verbose (bool) – If True, prints the test results.
datamodule (LightningDataModule | None) – A LightningDataModule that defines the test_dataloader hook.

For more information about multiple dataloaders, see this section.

Returns:

List of dictionaries with metrics logged during the test phase, e.g., in model- or callback hooks like test_step() etc. The length of the list corresponds to the number of test dataloaders used.

Raises:

TypeError – If no model is passed and there was no LightningModule passed in the previous run. If model passed is not LightningModule or torch._dynamo.OptimizedModule.
MisconfigurationException – If both dataloaders and datamodule are passed. Pass only one of these.
RuntimeError – If a compiled model is passed and the strategy is not supported.

Return type:

List[Mapping[str, float]]

validate(model=None, dataloaders=None, ckpt_path=None, verbose=True, datamodule=None)#

Perform one evaluation epoch over the validation set.

Parameters:

model (LightningModule | None) – The model to validate.
dataloaders (Any | LightningDataModule | None) – An iterable or collection of iterables specifying validation samples. Alternatively, a LightningDataModule that defines the val_dataloader hook.
ckpt_path (str | Path | None) – Either "best", "last", "hpc" or path to the checkpoint you wish to validate. If None and the model instance was passed, use the current weights. Otherwise, the best model checkpoint from the previous trainer.fit call will be loaded if a checkpoint callback is configured.
verbose (bool) – If True, prints the validation results.
datamodule (LightningDataModule | None) – A LightningDataModule that defines the val_dataloader hook.

For more information about multiple dataloaders, see this section.

Returns:

List of dictionaries with metrics logged during the validation phase, e.g., in model- or callback hooks like validation_step() etc. The length of the list corresponds to the number of validation dataloaders used.

Raises:

TypeError – If no model is passed and there was no LightningModule passed in the previous run. If model passed is not LightningModule or torch._dynamo.OptimizedModule.
MisconfigurationException – If both dataloaders and datamodule are passed. Pass only one of these.
RuntimeError – If a compiled model is passed and the strategy is not supported.

Return type:

List[Mapping[str, float]]

property accelerator: Accelerator#: !! processed by numpydoc !!

property callback_metrics: Dict[str, Tensor]#

The metrics available to callbacks.

def training_step(self, batch, batch_idx):
    self.log("a_val", 2.0)

callback_metrics = trainer.callback_metrics
assert callback_metrics["a_val"] == 2.0

property checkpoint_callback: Checkpoint | None#: The first ModelCheckpoint callback in the Trainer.callbacks list, or None if it doesn’t exist.

property checkpoint_callbacks: List[Checkpoint]#: A list of all instances of ModelCheckpoint found in the Trainer.callbacks list.

property ckpt_path: str | Path | None#

Set to the path/URL of a checkpoint loaded via fit(), validate(), test(), or predict().

None otherwise.

property current_epoch: int#: The current epoch, updated after the epoch end hooks are run.

property default_root_dir: str#

The default location to save artifacts of loggers, checkpoints etc.

It is used as a fallback if logger or checkpoint callback do not define specific save paths.

property device_ids: List[int]#: List of device indexes per node.

property distributed_sampler_kwargs: Dict[str, Any] | None#: !! processed by numpydoc !!

property early_stopping_callback: EarlyStopping | None#: The first EarlyStopping callback in the Trainer.callbacks list, or None if it doesn’t exist.

property early_stopping_callbacks: List[EarlyStopping]#: A list of all instances of EarlyStopping found in the Trainer.callbacks list.

property enable_validation: bool#: Check if we should run validation during training.

property estimated_stepping_batches: int | float#

The estimated number of batches that will optimizer.step() during training.

This accounts for gradient accumulation and the current trainer configuration. This might be used when setting up your training dataloader, if it hasn’t been set up already.

def configure_optimizers(self):
    optimizer = ...
    stepping_batches = self.trainer.estimated_stepping_batches
    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-3, total_steps=stepping_batches)
    return [optimizer], [scheduler]

Raises:: MisconfigurationException – If estimated stepping batches cannot be computed due to different accumulate_grad_batches at different epochs.

property evaluating: bool#: !! processed by numpydoc !!

property global_rank: int#: !! processed by numpydoc !!

property global_step: int#

The number of optimizer steps taken (does not reset each epoch).

This includes multiple optimizers (if enabled).

property interrupted: bool#: !! processed by numpydoc !!

property is_global_zero: bool#

Whether this process is the global zero in multi-node training.

def training_step(self, batch, batch_idx):
    if self.trainer.is_global_zero:
        print("in node 0, accelerator 0")

property is_last_batch: bool#: Whether trainer is executing the last batch.

property lightning_module: LightningModule#: !! processed by numpydoc !!

property local_rank: int#: !! processed by numpydoc !!

property log_dir: str | None#

The directory for the current experiment. Use this to save images to, etc…

Note

You must call this on all processes. Failing to do so will cause your program to stall forever.

def training_step(self, batch, batch_idx):
    img = ...
    save_img(img, self.trainer.log_dir)

property logged_metrics: Dict[str, Tensor]#

The metrics sent to the loggers.

This includes metrics logged via log() with the :paramref:`~lightning.pytorch.core.LightningModule.log.logger` argument set.

property logger: Logger | None#: The first Logger being used.

property loggers: List[Logger]#

The list of Logger used.

for logger in trainer.loggers:
    logger.log_metrics({"foo": 1.0})

property lr_scheduler_configs: List[LRSchedulerConfig]#: !! processed by numpydoc !!

property max_epochs: int | None#: !! processed by numpydoc !!

property max_steps: int#: !! processed by numpydoc !!

property min_epochs: int | None#: !! processed by numpydoc !!

property min_steps: int | None#: !! processed by numpydoc !!

property model: Module | None#

The LightningModule, but possibly wrapped into DataParallel or DistributedDataParallel.

To access the pure LightningModule, use lightning_module() instead.

property node_rank: int#: !! processed by numpydoc !!

property num_devices: int#: Number of devices the trainer uses per node.

property num_nodes: int#: !! processed by numpydoc !!

property num_predict_batches: List[int | float]#: The number of prediction batches that will be used during trainer.predict().

property num_sanity_val_batches: List[int | float]#: The number of validation batches that will be used during the sanity-checking part of trainer.fit().

property num_test_batches: List[int | float]#: The number of test batches that will be used during trainer.test().

property num_training_batches: int | float#: The number of training batches that will be used during trainer.fit().

property num_val_batches: List[int | float]#: The number of validation batches that will be used during trainer.fit() or trainer.validate().

property optimizers: List[Optimizer]#: !! processed by numpydoc !!

property precision: Literal['transformer-engine', 'transformer-engine-float16', '16-true', '16-mixed', 'bf16-true', 'bf16-mixed', '32-true', '64-true']#: !! processed by numpydoc !!

property precision_plugin: Precision#: !! processed by numpydoc !!

property predict_dataloaders: Any | None#: The prediction dataloader(s) used during trainer.predict().

property predicting: bool#: !! processed by numpydoc !!

property progress_bar_callback: ProgressBar | None#: An instance of ProgressBar found in the Trainer.callbacks list, or None if one doesn’t exist.

property progress_bar_metrics: Dict[str, float]#

The metrics sent to the progress bar.

This includes metrics logged via log() with the :paramref:`~lightning.pytorch.core.LightningModule.log.prog_bar` argument set.

property received_sigterm: bool#

Whether a signal.SIGTERM signal was received.

For example, this can be checked to exit gracefully.

property sanity_checking: bool#

Whether sanity checking is running.

Useful to disable some hooks, logging or callbacks during the sanity checking.

property scaler: Any | None#: !! processed by numpydoc !!

property strategy: Strategy#: !! processed by numpydoc !!

property test_dataloaders: Any | None#: The test dataloader(s) used during trainer.test().

property testing: bool#: !! processed by numpydoc !!

property train_dataloader: Any | None#: The training dataloader(s) used during trainer.fit().

property training: bool#: !! processed by numpydoc !!

property val_dataloaders: Any | None#: The validation dataloader(s) used during trainer.fit() or trainer.validate().

property validating: bool#: !! processed by numpydoc !!

property world_size: int#: !! processed by numpydoc !!

count_number_of_parameters(model, only_trainable=True)#

Count the number of trainable params.

If all params, specify only_trainable = False.

Ref:

Parameters:

modeltorch.nn.Module: The model.
only_trainablebool, optional: If True, only count trainable parameters (default: True).

Returns:

int: The number of parameters.

extras(cfg)#

Apply optional utilities before the task is started.

Utilities:

Ignoring python warnings.
Setting tags from command line.
Rich config printing.

Parameters:

cfgDictConfig: A DictConfig object containing the config tree.

get_default_metrics(task, metrics=None)#

Get default metrics for a given task.

Parameters:

taskstr: Task, either “classification” or “regression”.
metricslist, optional: List of metrics to be used. If None, the default metrics will be used.

Returns:

list: List of default metrics.

Raises:

ValueError: If the task is invalid.

get_default_trainer()#

Get default trainer configuration.

Returns:

str: Default trainer configuration file name.

get_default_transform(dataset, model)#

Get default transform for a given data domain and model.

Parameters:

datasetstr: Dataset name. Should be in the format “data_domain/name”.
modelstr: Model name. Should be in the format “model_domain/name”.

Returns:

str: Default transform.

get_flattened_channels(num_nodes, channels)#

Get the output dimension of flattening a feature matrix.

Parameters:

num_nodesint: Hidden dimension for the first layer.
channelsint: Channel dimension.

Returns:

int: Flatenned cchannels dimension.

get_metric_value(metric_dict, metric_name)#

Safely retrieves value of the metric logged in LightningModule.

Parameters:

metric_dictdict: A dict containing metric values.
metric_namestr, optional: If provided, the name of the metric to retrieve.

Returns:

float, None: If a metric name was provided, the value of the metric.

get_monitor_metric(task, metric)#

Get monitor metric for a given task.

Parameters:

taskstr: Task, either “classification” or “regression”.
metricstr: Name of the metric function.

Returns:

str: Monitor metric.

Raises:

ValueError: If the task is invalid.

get_monitor_mode(task)#

Get monitor mode for a given task.

Parameters:

taskstr: Task, either “classification” or “regression”.

Returns:

str: Monitor mode, either “max” or “min”.

Raises:

ValueError: If the task is invalid.

get_non_relational_out_channels(num_nodes, channels, task_level)#

Get the output dimension for a non-relational model.

Parameters:

num_nodesint: Number of nodes in the input graph.
channelsint: Channel dimension.
task_levelint: Task level for the model.

Returns:

int: Output dimension.

get_required_lifting(data_domain, model)#

Get required transform for a given data domain and model.

Parameters:

data_domainstr: Dataset domain.
modelstr: Model name. Should be in the format “model_domain/name”.

Returns:

str: Required transform.

infer_in_channels(dataset, transforms)#

Infer the number of input channels for a given dataset.

Parameters:

datasetDictConfig: Configuration parameters for the dataset.
transformsDictConfig: Configuration parameters for the transforms.

Returns:

list: List with dimensions of the input channels.

infer_num_cell_dimensions(selected_dimensions, in_channels)#

Infer the length of a list.

Parameters:

selected_dimensionslist: List of selected dimensions. If not None it will be used to infer the length.
in_channelslist: List of input channels. If selected_dimensions is None, this list will be used to infer the length.

Returns:

int: Length of the input list.

infer_topotune_num_cell_dimensions(neighborhoods)#

Infer the length of a list.

Parameters:

neighborhoodslist: List of neighborhoods.

Returns:

int: Length of the input list.

initialize_hydra()#

Initialize Hydra when main is not an option (e.g. tests).

Returns:

DictConfig: A DictConfig object containing the config tree.

instantiate_callbacks(callbacks_cfg)#

Instantiate callbacks from config.

Parameters:

callbacks_cfgDictConfig: A DictConfig object containing callback configurations.

Returns:

list[Callback]: A list of instantiated callbacks.

instantiate_loggers(logger_cfg)#

Instantiate loggers from config.

Parameters:

logger_cfgDictConfig: A DictConfig object containing logger configurations.

Returns:

list[Logger]: A list of instantiated loggers.

log_hyperparameters(object_dict)#

Control which config parts are saved by Lightning loggers.

Additionally saves:

Number of model parameters

Parameters:

object_dictdict[str, Any]

A dictionary containing the following objects:

“cfg”: A DictConfig object containing the main config.
“model”: The Lightning model.
“trainer”: The Lightning trainer.

main(cfg)#

Main entry point for training.

Parameters:

cfgDictConfig: Configuration composed by Hydra.

Returns:

float | None: Optional[float] with optimized metric value.

run(cfg)#

Wrapper function that executes the task function.

Parameters:

cfgDictConfig: A DictConfig object containing the config tree.

Returns:

tuple[dict[str, Any], dict[str, Any]]: The metric and object dictionaries returned by the task function.

task_wrapper(task_func)#

Optional decorator that controls the failure behavior when executing the task function.

This wrapper can be used to: - make sure loggers are closed even if the task function raises an exception (prevents multirun failure). - save the exception to a .log file. - mark the run as failed with a dedicated file in the logs/ folder (so we can find and rerun it later). - etc. (adjust depending on your needs).

Example: ``` @utils.task_wrapper def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:

… return metric_dict, object_dict

```

Parameters:

task_funcCallable: The task function to be wrapped.

Returns:

Callable: The wrapped task function.