ml.trainers.mixins.profiler

Defines a trainer mixin for profiling PyTorch models.

This can be enabled by setting trainer.profiler.enabled=true in your configuration file. It will run the first few steps of training with the profiler enabled, and then log a report of the results to the log directory.

class ml.trainers.mixins.profiler.Profiler(enabled: bool = False, record_shapes: bool = False, profile_memory: bool = False, with_stack: bool = False, with_flops: bool = False, with_modules: bool = False, wait_steps: int = 10, warmup_steps: int = 10, active_steps: int = 10, repeat_steps: int = 1, skip_first_steps: int = 10, table_size: int = 10)

Bases: object

enabled: bool = False

record_shapes: bool = False

profile_memory: bool = False

with_stack: bool = False

with_flops: bool = False

with_modules: bool = False

wait_steps: int = 10

warmup_steps: int = 10

active_steps: int = 10

repeat_steps: int = 1

skip_first_steps: int = 10

table_size: int = 10

class ml.trainers.mixins.profiler.ProfilerTrainerConfig(name: str = '???', exp_name: str = '${ml.exp_name:null}', exp_dir: str = '???', log_dir_name: str = 'logs', use_double_weight_precision: bool = False, checkpoint: ml.trainers.base.CheckpointConfig = <factory>, profiler: ml.trainers.mixins.profiler.Profiler = <factory>)

Bases: BaseTrainerConfig

profiler: Profiler

class ml.trainers.mixins.profiler.ProfilerTrainerMixin(config: ProfilerTrainerConfigT)

Bases: StepContextMixin[ProfilerTrainerConfigT, ModelT, TaskT], BaseTrainer[ProfilerTrainerConfigT, ModelT, TaskT]

Defines a trainer mixin for enabling the PyTorch profiler.

step_context(step: Literal['backward', 'build_rl_dataset', 'change_mode', 'clip_grads', 'collect_rl_samples', 'forward', 'get_single_loss', 'log_losses', 'on_epoch_end', 'on_epoch_start', 'on_step_end', 'on_step_start', 'step', 'update_state', 'write_logs', 'zero_grads']) → ContextManager

write_logs(task: TaskT, model: ModelT, state: State) → None

on_profiler_trace_ready(prof: profile) → None

get_profile() → profile | None