PyTorchProfiler

class torchtnt.framework.callbacks.PyTorchProfiler(profiler: profile)

A callback which profiles user code using PyTorch Profiler.

Parameters:profiler – a torch.profiler.profile context manager which will be used

on_eval_end(state: State, unit: EvalUnit[TEvalData]) → None

Hook called after evaluation ends.

on_eval_start(state: State, unit: EvalUnit[TEvalData]) → None

Hook called before evaluation starts.

on_eval_step_end(state: State, unit: EvalUnit[TEvalData]) → None

Hook called after an eval step ends.

on_predict_end(state: State, unit: PredictUnit[TPredictData]) → None

Hook called after prediction ends.

on_predict_start(state: State, unit: PredictUnit[TPredictData]) → None

Hook called before prediction starts.

on_predict_step_end(state: State, unit: PredictUnit[TPredictData]) → None

Hook called after a predict step ends.

on_train_end(state: State, unit: TrainUnit[TTrainData]) → None

Hook called after training ends.

on_train_start(state: State, unit: TrainUnit[TTrainData]) → None

Hook called before training starts.

on_train_step_end(state: State, unit: TrainUnit[TTrainData]) → None

Hook called after a train step ends.