aitune.torch.tune_strategy.mixin.find_max_batch_size_mixin

Find max batch size mixin for tune strategy.

Looks for best batch size for the module using Torch Eager backend.

Module Contents

Classes

Name	Description
`FindMaxBatchSizeMixin`	TuneStrategy mixin that finds max batch size.

API

class aitune.torch.tune_strategy.mixin.find_max_batch_size_mixin.FindMaxBatchSizeMixin(
    args = (),
    kwargs = {}
)

Bases: TuneStrategy

TuneStrategy mixin that finds max batch size.

_find_max_batch_size_backend_class

type[Backend] = TorchEagerBackend

aitune.torch.tune_strategy.mixin.find_max_batch_size_mixin.FindMaxBatchSizeMixin._pre_tune(
    module: torch.nn.Module,
    name: str,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    device: torch.device,
    cache_dir: pathlib.Path
)

Extends tune method to find max batch size.

aitune.torch.tune_strategy.mixin.find_max_batch_size_mixin.FindMaxBatchSizeMixin.enable_find_max_batch_size(
    enable: bool = True
) -> aitune.torch.tune_strategy.mixin.find_max_batch_size_mixin.FindMaxBatchSizeMixin

Enables or disables find max batch size.

aitune.torch.tune_strategy.mixin.find_max_batch_size_mixin.FindMaxBatchSizeMixin.find_max_batch_size(
    module: torch.nn.Module,
    name: str,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    device: torch.device,
    cache_dir: pathlib.Path
)

Finds max batch size for the module.

aitune.torch.tune_strategy.mixin.find_max_batch_size_mixin.FindMaxBatchSizeMixin.set_find_max_batch_size_default_backend_class(
    default_backend_class: type[aitune.torch.backend.backend.Backend]
) -> aitune.torch.tune_strategy.mixin.find_max_batch_size_mixin.FindMaxBatchSizeMixin

Sets default backend class for find max batch size.