modlee.model.callbacks module

class modlee.model.callbacks.DataMetafeaturesCallback(data_snapshot_size=10000000.0, DataMetafeatures=None, *args, **kwargs)

Bases: ModleeCallback

Callback to calculate and log data meta-features.

on_train_start(trainer: Trainer, pl_module: LightningModule) → None

Called when the train begins.

class modlee.model.callbacks.LogCodeTextCallback(kwargs_to_cache={}, *args, **kwargs)

Bases: ModleeCallback

Callback to log the model as code and text.

on_train_start(trainer: Trainer, pl_module: LightningModule) → None

Called when the train begins.

setup(trainer: Trainer, pl_module: LightningModule, stage: str) → None

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

class modlee.model.callbacks.LogONNXCallback

Bases: ModleeCallback

Callback for logging the model in its ONNX representations. Deprecated, will be combined with LogCodeTextCallback.

on_train_start(trainer: Trainer, pl_module: LightningModule) → None

Called when the train begins.

setup(trainer: Trainer, pl_module: LightningModule, stage: str) → None

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

class modlee.model.callbacks.LogOutputCallback(*args, **kwargs)

Bases: Callback

Callback to log the output metrics for each batch.

class modlee.model.callbacks.LogParamsCallback

Bases: Callback

Callback to log parameters at the start of training.

on_train_start(trainer: Trainer, pl_module: LightningModule) → None

Called when the train begins.

class modlee.model.callbacks.LogTransformsCallback

Bases: ModleeCallback

Logs transforms applied to the dataset, if applied with torchvision.transforms

on_train_start(trainer, pl_module)

Called when the train begins.

class modlee.model.callbacks.ModleeCallback

Bases: Callback

Base class for Modlee-specific callbacks.

get_input(trainer, pl_module)

Get an input (one element from a batch) from a trainer’s dataloader.

Parameters:

• trainer – The trainer with the dataloader.

• pl_module – The model module, used for loading the data input to the correct device.

Returns:

An input from the batch.

class modlee.model.callbacks.PushServerCallback

Bases: Callback

Callback to push run assets to the server at the end of training.

on_fit_end(trainer: Trainer, pl_module: LightningModule) → None

Called when fit ends.