otx.backend.native.callbacks#

Module for OTX custom callbacks.

Classes

BatchSizeFinder([steps_per_trial])

This callback makes trainer run specified iteration and exit.

class otx.backend.native.callbacks.BatchSizeFinder(steps_per_trial: int = 3)[source]#

Bases: Callback

This callback makes trainer run specified iteration and exit.

Parameters:: steps_per_trial – number of steps to run with a given batch size. Ideally 1 should be enough to test if an OOM error occurs, however in practice a few are needed.

on_fit_start(trainer: Trainer, pl_module: LightningModule) → None[source]#: Run steps_per_trial iterations and exit.

setup(trainer: Trainer, pl_module: LightningModule, stage: str | None = None) → None[source]#: Check current stage is fit.