Skip to content

Trainers

TabularTrainer

TabularTrainer(
    model: TabularModel, dp_budget: DpBudget | None = None
)

Trainer for a TabularModel.

Parameters:

Name Type Description Default
model TabularModel

A TabularModel instance.

required
dp_budget DpBudget | None

The (eps, delta)-budget for differentially private (DP) training. If None (the default), the training will not be differentially private. Available only for single table datasets.

None

train

train(
    dataset: TabularDataset,
    batch_size: int | AutoBatch,
    n_epochs: int | None = None,
    n_steps: int | None = None,
    lr: float = 0.0,
    valid: Validation | None = None,
    hooks: Sequence[TrainHook] = (),
    accumulate_grad: int = 1,
    dp_step: DpStep | None = None,
    world_size: int = 0,
) -> None

Train the TabularModel with the input TabularDataset.

Parameters:

Name Type Description Default
dataset TabularDataset

The training data, as a TabularDataset object.

required
batch_size int | AutoBatch

The size of a batch of data during training. It is possible to specify an instance of AutoBatch to automatically estimate the optimal batch size.

required
n_epochs int | None

The number of training epochs. One and only one of n_epochs and n_steps must be provided.

None
n_steps int | None

The number of training steps. One and only one of n_epochs and n_steps must be provided.

None
lr float

The learning rate. If it is 0 the optimal value for the learning rate is automatically determined.

0.0
valid Validation | None

A Validation object. If None, no validation is performed.

None
hooks Sequence[TrainHook]

A sequence of custom TrainHook objects.

()
accumulate_grad int

The number of gradient accumulation steps. If equal to 1, the weights are updated at each step.

1
dp_step DpStep | None

Data for differentially private step. Must be provided if and only if the trainer has a DP-budget.

None
world_size int

Number of GPUs where to distribute the training. If 0, the training is performed on a single device, on the current device of the TabularTrainer object.

0

save

save(path: Path | str) -> None

Save the TabularTrainer to a checkpoint at the given path.

Parameters:

Name Type Description Default
path Path | str

The path where to save the checkpoint.

required

load classmethod

load(path: Path | str) -> TabularTrainer

Load the TabularTrainer from the checkpoint at the given path.

Parameters:

Name Type Description Default
path Path | str

The path to the loaded checkpoint.

required

Returns:

Type Description
TabularTrainer

The loaded TabularTrainer.

TextTrainer

TextTrainer(model: TextModel)

Trainer for a TextModel.

Parameters:

Name Type Description Default
model TextModel

A TextModel instance.

required

train

train(
    dataset: TextDataset,
    batch_size: int | AutoBatch,
    n_epochs: int | None = None,
    n_steps: int | None = None,
    lr: float = 0.0,
    valid: Validation | None = None,
    hooks: Sequence[TrainHook] = (),
    accumulate_grad: int = 1,
    world_size: int = 0,
) -> None

Train the TextModel with the input TextDataset.

Parameters:

Name Type Description Default
dataset TextDataset

The training data, as a TextDataset object.

required
batch_size int | AutoBatch

The size of a batch of data during training. It is possible to specify an instance of AutoBatch to automatically estimate the optimal batch size.

required
n_epochs int | None

The number of training epochs. One and only one of n_epochs and n_steps must be provided.

None
n_steps int | None

The number of training steps. One and only one of n_epochs and n_steps must be provided.

None
lr float

The learning rate. If it is 0 the optimal value for the learning rate is automatically determined.

0.0
valid Validation | None

A Validation object. If None, no validation is performed.

None
hooks Sequence[TrainHook]

A sequence of custom TrainHook objects.

()
accumulate_grad int

The number of gradient accumulation steps. If equal to 1, the weights are updated at each step.

1
world_size int

Number of GPUs where to distribute the training. If 0, the training is performed on a single device, on the current device of the TextTrainer object.

0

save

save(path: Path | str) -> None

Save the TextTrainer to a checkpoint at the given path.

Parameters:

Name Type Description Default
path Path | str

The path where to save the checkpoint.

required

load classmethod

load(path: Path | str) -> TextTrainer

Load the TextTrainer from the checkpoint at the given path.

Parameters:

Name Type Description Default
path Path | str

The path to the loaded checkpoint.

required

Returns:

Type Description
TextTrainer

The loaded TextTrainer.

AutoBatch pydantic-config

AutoBatch(
    memory: PositiveInt,
    min_batch_size: PositiveInt = 1,
    max_batch_size: PositiveInt = 256,
    increment_factor: Annotated[float, Field(gt=1)] = sqrt(
        2
    ),
    min_steps_epoch: Annotated[int, Field(ge=1)] = 50,
)

Configuration for automatic batch size estimation, given the available (CPU or GPU) memory. The optimal batch size is obtained by estimating the memory required for a training step, for increasing values of the batch size. The selected batch size is the largest that fits within the provided memory constraints.

Parameters:

Name Type Description Default
memory PositiveInt

The available memory in MB.

required
min_batch_size PositiveInt

The minimum batch size.

1
max_batch_size PositiveInt

The maximum batch size.

256
increment_factor Annotated[float, Field(gt=1)]

The multiplicative factor used to compute the increasing batch sizes.

sqrt(2)
min_steps_epoch Annotated[int, Field(ge=1)]

The minimum number of steps per epoch to ensure when computing the maximum batch size tested.

50