multimeditron package

Subpackages

Submodules

multimeditron.profiling module

class multimeditron.profiling.NvtxAnnotationCallback(global_step_start=100, global_step_stop=120)

Bases: TrainerCallback

” Adding NVTX annotations for profiling with Nsight Systems.

on_epoch_begin(args, state, control, **kwargs)

Event called at the beginning of an epoch.

on_epoch_end(args, state, control, **kwargs)

Event called at the end of an epoch.

on_evaluate(args, state, control, output_metrics, **kwargs)

Event called after an evaluation phase.

on_forward_begin(args, state, control, **kwargs)
on_forward_end(args, state, control, **kwargs)
on_init_end(args, state, control, **kwargs)

Event called at the end of the initialization of the [Trainer].

on_log(args, state, control, logs, **kwargs)

Event called after logging the last logs.

on_optimizer_step(args, state, control, **kwargs)

Event called after the optimizer step but before gradients are zeroed out. Useful for monitoring gradients.

on_pre_optimizer_step(args, state, control, **kwargs)

Event called before the optimizer step but after gradient clipping. Useful for monitoring gradients.

on_predict(args, state, control, output_metrics, **kwargs)

Event called after a successful prediction.

on_prediction_step(args, state, control, **kwargs)

Event called after a prediction step.

on_prepare_inputs_begin(args, state, control, **kwargs)
on_prepare_inputs_end(args, state, control, **kwargs)
on_save(args, state, control, **kwargs)

Event called after a checkpoint save.

on_step_begin(args, state, control, **kwargs)

Event called at the beginning of a training step. If using gradient accumulation, one training step might take several inputs.

on_step_end(args, state, control, **kwargs)

Event called at the end of a training step. If using gradient accumulation, one training step might take several inputs.

on_substep_end(args, state, control, **kwargs)

Event called at the end of an substep during gradient accumulation.

on_train_begin(args, state, control, **kwargs)

Event called at the beginning of training.

on_train_end(args, state, control, **kwargs)

Event called at the end of training.

Module contents