Shortcuts

Performance and Bottleneck Profiler

Profiling your training run can help you understand if there are any bottlenecks in your code.

Built-in checks

PyTorch Lightning supports profiling standard actions in the training loop out of the box, including:

  • on_epoch_start

  • on_epoch_end

  • on_batch_start

  • tbptt_split_batch

  • model_forward

  • model_backward

  • on_after_backward

  • optimizer_step

  • on_batch_end

  • training_step_end

  • on_training_end

Enable simple profiling

If you only wish to profile the standard actions, you can set profiler=True when constructing your Trainer object.

trainer = Trainer(..., profiler=True)

The profiler’s results will be printed at the completion of a training fit().

Profiler Report

Action                  |  Mean duration (s)    |  Total time (s)
-----------------------------------------------------------------
on_epoch_start          |  5.993e-06            |  5.993e-06
get_train_batch         |  0.0087412            |  16.398
on_batch_start          |  5.0865e-06           |  0.0095372
model_forward           |  0.0017818            |  3.3408
model_backward          |  0.0018283            |  3.4282
on_after_backward       |  4.2862e-06           |  0.0080366
optimizer_step          |  0.0011072            |  2.0759
on_batch_end            |  4.5202e-06           |  0.0084753
on_epoch_end            |  3.919e-06            |  3.919e-06
on_train_end            |  5.449e-06            |  5.449e-06

Advanced Profiling

If you want more information on the functions called during each event, you can use the AdvancedProfiler. This option uses Python’s cProfiler to provide a report of time spent on each function called within your code.

profiler = AdvancedProfiler()
trainer = Trainer(..., profiler=profiler)

The profiler’s results will be printed at the completion of a training fit(). This profiler report can be quite long, so you can also specify an output_filename to save the report instead of logging it to the output in your terminal. The output below shows the profiling for the action get_train_batch.

Profiler Report

Profile stats for: get_train_batch
        4869394 function calls (4863767 primitive calls) in 18.893 seconds
Ordered by: cumulative time
List reduced from 76 to 10 due to restriction <10>
ncalls  tottime  percall  cumtime  percall filename:lineno(function)
3752/1876    0.011    0.000   18.887    0.010 {built-in method builtins.next}
    1876     0.008    0.000   18.877    0.010 dataloader.py:344(__next__)
    1876     0.074    0.000   18.869    0.010 dataloader.py:383(_next_data)
    1875     0.012    0.000   18.721    0.010 fetch.py:42(fetch)
    1875     0.084    0.000   18.290    0.010 fetch.py:44(<listcomp>)
    60000    1.759    0.000   18.206    0.000 mnist.py:80(__getitem__)
    60000    0.267    0.000   13.022    0.000 transforms.py:68(__call__)
    60000    0.182    0.000    7.020    0.000 transforms.py:93(__call__)
    60000    1.651    0.000    6.839    0.000 functional.py:42(to_tensor)
    60000    0.260    0.000    5.734    0.000 transforms.py:167(__call__)

You can also reference this profiler in your LightningModule to profile specific actions of interest. If you don’t want to always have the profiler turned on, you can optionally pass a PassThroughProfiler which will allow you to skip profiling without having to make any code changes. Each profiler has a method profile() which returns a context handler. Simply pass in the name of your action that you want to track and the profiler will record performance for code executed within this context.

from pytorch_lightning.profiler import Profiler, PassThroughProfiler

class MyModel(LightningModule):
    def __init__(self, profiler=None):
        self.profiler = profiler or PassThroughProfiler()

    def custom_processing_step(self, data):
        with profiler.profile('my_custom_action'):
            # custom processing step
        return data

profiler = Profiler()
model = MyModel(profiler)
trainer = Trainer(profiler=profiler, max_epochs=1)
class pytorch_lightning.profiler.BaseProfiler(output_streams=None)[source]

Bases: abc.ABC

If you wish to write a custom profiler, you should inhereit from this class.

Params:

stream_out: callable

describe()[source]

Logs a profile report after the conclusion of the training run.

Return type

None

profile(action_name)[source]

Yields a context manager to encapsulate the scope of a profiled action.

Example:

with self.profile('load training data'):
    # load training data code

The profiler will start once you’ve entered the context and will automatically stop once you exit the code block.

Return type

None

abstract start(action_name)[source]

Defines how to start recording an action.

Return type

None

abstract stop(action_name)[source]

Defines how to record the duration once an action is complete.

Return type

None

abstract summary()[source]

Create profiler summary in text format.

Return type

str

class pytorch_lightning.profiler.SimpleProfiler(output_filename=None)[source]

Bases: pytorch_lightning.profiler.profilers.BaseProfiler

This profiler simply records the duration of actions (in seconds) and reports the mean duration of each action and the total time spent over the entire training run.

Params:
output_filename (str): optionally save profile results to file instead of printing

to std out when training is finished.

describe()[source]

Logs a profile report after the conclusion of the training run.

start(action_name)[source]

Defines how to start recording an action.

Return type

None

stop(action_name)[source]

Defines how to record the duration once an action is complete.

Return type

None

summary()[source]

Create profiler summary in text format.

Return type

str

class pytorch_lightning.profiler.AdvancedProfiler(output_filename=None, line_count_restriction=1.0)[source]

Bases: pytorch_lightning.profiler.profilers.BaseProfiler

This profiler uses Python’s cProfiler to record more detailed information about time spent in each function call recorded during a given action. The output is quite verbose and you should only use this if you want very detailed reports.

Parameters
  • output_filename (Optional[str]) – optionally save profile results to file instead of printing to std out when training is finished.

  • line_count_restriction (float) – this can be used to limit the number of functions reported for each action. either an integer (to select a count of lines), or a decimal fraction between 0.0 and 1.0 inclusive (to select a percentage of lines)

describe()[source]

Logs a profile report after the conclusion of the training run.

start(action_name)[source]

Defines how to start recording an action.

Return type

None

stop(action_name)[source]

Defines how to record the duration once an action is complete.

Return type

None

summary()[source]

Create profiler summary in text format.

Return type

str

class pytorch_lightning.profiler.PassThroughProfiler[source]

Bases: pytorch_lightning.profiler.profilers.BaseProfiler

This class should be used when you don’t want the (small) overhead of profiling. The Trainer uses this class by default.

Params: stream_out: callable

start(action_name)[source]

Defines how to start recording an action.

Return type

None

stop(action_name)[source]

Defines how to record the duration once an action is complete.

Return type

None

summary()[source]

Create profiler summary in text format.

Return type

str