LightningIRTrainer

class lightning_ir.main.LightningIRTrainer(*, accelerator: str | Accelerator = 'auto', strategy: str | Strategy = 'auto', devices: list[int] | str | int = 'auto', num_nodes: int = 1, precision: Literal[64, 32, 16] | Literal['transformer-engine', 'transformer-engine-float16', '16-true', '16-mixed', 'bf16-true', 'bf16-mixed', '32-true', '64-true'] | Literal['64', '32', '16', 'bf16'] | None = None, logger: Logger | Iterable[Logger] | bool | None = None, callbacks: list[Callback] | Callback | None = None, fast_dev_run: int | bool = False, max_epochs: int | None = None, min_epochs: int | None = None, max_steps: int = -1, min_steps: int | None = None, max_time: str | timedelta | dict[str, int] | None = None, limit_train_batches: int | float | None = None, limit_val_batches: int | float | None = None, limit_test_batches: int | float | None = None, limit_predict_batches: int | float | None = None, overfit_batches: int | float = 0.0, val_check_interval: int | float | None = None, check_val_every_n_epoch: int | None = 1, num_sanity_val_steps: int | None = None, log_every_n_steps: int | None = None, enable_checkpointing: bool | None = None, enable_progress_bar: bool | None = None, enable_model_summary: bool | None = None, accumulate_grad_batches: int = 1, gradient_clip_val: int | float | None = None, gradient_clip_algorithm: str | None = None, deterministic: bool | Literal['warn'] | None = None, benchmark: bool | None = None, inference_mode: bool = True, use_distributed_sampler: bool = True, profiler: Profiler | str | None = None, detect_anomaly: bool = False, barebones: bool = False, plugins: Precision | ClusterEnvironment | CheckpointIO | LayerSync | list[Precision | ClusterEnvironment | CheckpointIO | LayerSync] | None = None, sync_batchnorm: bool = False, reload_dataloaders_every_n_epochs: int = 0, default_root_dir: str | Path | None = None, model_registry: str | None = None)[source]

Bases: Trainer

__init__(*, accelerator: str | Accelerator = 'auto', strategy: str | Strategy = 'auto', devices: list[int] | str | int = 'auto', num_nodes: int = 1, precision: Literal[64, 32, 16] | Literal['transformer-engine', 'transformer-engine-float16', '16-true', '16-mixed', 'bf16-true', 'bf16-mixed', '32-true', '64-true'] | Literal['64', '32', '16', 'bf16'] | None = None, logger: Logger | Iterable[Logger] | bool | None = None, callbacks: list[Callback] | Callback | None = None, fast_dev_run: int | bool = False, max_epochs: int | None = None, min_epochs: int | None = None, max_steps: int = -1, min_steps: int | None = None, max_time: str | timedelta | dict[str, int] | None = None, limit_train_batches: int | float | None = None, limit_val_batches: int | float | None = None, limit_test_batches: int | float | None = None, limit_predict_batches: int | float | None = None, overfit_batches: int | float = 0.0, val_check_interval: int | float | None = None, check_val_every_n_epoch: int | None = 1, num_sanity_val_steps: int | None = None, log_every_n_steps: int | None = None, enable_checkpointing: bool | None = None, enable_progress_bar: bool | None = None, enable_model_summary: bool | None = None, accumulate_grad_batches: int = 1, gradient_clip_val: int | float | None = None, gradient_clip_algorithm: str | None = None, deterministic: bool | Literal['warn'] | None = None, benchmark: bool | None = None, inference_mode: bool = True, use_distributed_sampler: bool = True, profiler: Profiler | str | None = None, detect_anomaly: bool = False, barebones: bool = False, plugins: Precision | ClusterEnvironment | CheckpointIO | LayerSync | list[Precision | ClusterEnvironment | CheckpointIO | LayerSync] | None = None, sync_batchnorm: bool = False, reload_dataloaders_every_n_epochs: int = 0, default_root_dir: str | Path | None = None, model_registry: str | None = None) → None

Customize every aspect of training via flags.

Parameters:

  • accelerator – Supports passing different accelerator types (“cpu”, “gpu”, “tpu”, “hpu”, “mps”, “auto”) as well as custom accelerator instances.

  • strategy – Supports different training strategies with aliases as well custom strategies. Default: "auto".

  • devices – 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 – Number of GPU nodes for distributed training. Default: 1.

  • precision – 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 (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 – Add a callback or list of callbacks. Default: None.

  • fast_dev_run – 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 – 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 – Force training for at least these many epochs. Disabled by default (None).

  • max_steps – 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 – Force training for at least these number of steps. Disabled by default (None).

  • max_time – 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 – How much of training dataset to check (float = fraction, int = num_batches). Default: 1.0.

  • limit_val_batches – How much of validation dataset to check (float = fraction, int = num_batches). Default: 1.0.

  • limit_test_batches – How much of test dataset to check (float = fraction, int = num_batches). Default: 1.0.

  • limit_predict_batches – How much of prediction dataset to check (float = fraction, int = num_batches). Default: 1.0.

  • overfit_batches – Overfit a fraction of training/validation data (float) or a set number of batches (int). Default: 0.0.

  • val_check_interval – 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 – 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 – 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 – How often to log within steps. Default: 50.

  • enable_checkpointing – 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 – Whether to enable to progress bar by default. Default: True.

  • enable_model_summary – Whether to enable model summarization by default. Default: True.

  • accumulate_grad_batches – Accumulates gradients over k batches before stepping the optimizer. Default: 1.

  • gradient_clip_val – 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 – 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 – 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 – 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 – Whether to use torch.inference_mode() or torch.no_grad() during evaluation (validate/test/predict).

  • use_distributed_sampler – 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 – To profile individual steps during training and assist in identifying bottlenecks. Default: None.

  • detect_anomaly – Enable anomaly detection for the autograd engine. Default: False.

  • barebones – 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 – Plugins allow modification of core behavior like ddp and amp, and enable custom lightning plugins. Default: None.

  • sync_batchnorm – Synchronize batch norm layers between process groups/whole world. Default: False.

  • reload_dataloaders_every_n_epochs – Set to a positive integer to reload dataloaders every n epochs. Default: 0.

  • default_root_dir – 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/’

  • model_registry – The name of the model being uploaded to Model hub.

Raises:

  • TypeError – If gradient_clip_val is not an int or float.

  • MisconfigurationException – If gradient_clip_algorithm is invalid.

Methods

index([model, dataloaders, ckpt_path, ...])

Index a collection of documents.

re_rank([model, dataloaders, ckpt_path, ...])

Re-rank a set of retrieved documents.

search([model, dataloaders, ckpt_path, ...])

Search for relevant documents.

Attributes

index(model: LightningModule | None = None, dataloaders: Any | LightningDataModule | None = None, ckpt_path: str | Path | None = None, verbose: bool = True, datamodule: LightningDataModule | None = None) → List[Mapping[str, float]][source]

Index a collection of documents.

re_rank(model: LightningModule | None = None, dataloaders: Any | LightningDataModule | None = None, ckpt_path: str | Path | None = None, verbose: bool = True, datamodule: LightningDataModule | None = None) → List[Mapping[str, float]][source]

Re-rank a set of retrieved documents.

search(model: LightningModule | None = None, dataloaders: Any | LightningDataModule | None = None, ckpt_path: str | Path | None = None, verbose: bool = True, datamodule: LightningDataModule | None = None) → List[Mapping[str, float]][source]

Search for relevant documents.