Shortcuts

Find bottlenecks in your code (expert)

Audience: Users who want to build their own profilers.


Build your own profiler

To build your own profiler, subclass Profiler and override some of its methods. Here is a simple example that profiles the first occurrence and total calls of each action:

from pytorch_lightning.profilers import Profiler
from collections import defaultdict
import time


class ActionCountProfiler(Profiler):
    def __init__(self, dirpath=None, filename=None):
        super().__init__(dirpath=dirpath, filename=filename)
        self._action_count = defaultdict(int)
        self._action_first_occurrence = {}

    def start(self, action_name):
        if action_name not in self._action_first_occurrence:
            self._action_first_occurrence[action_name] = time.strftime("%m/%d/%Y, %H:%M:%S")

    def stop(self, action_name):
        self._action_count[action_name] += 1

    def summary(self):
        res = f"\nProfile Summary: \n"
        max_len = max(len(x) for x in self._action_count)

        for action_name in self._action_count:
            # generate summary for actions called more than once
            if self._action_count[action_name] > 1:
                res += (
                    f"{action_name:<{max_len}s} \t "
                    + "self._action_first_occurrence[action_name]} \t "
                    + "{self._action_count[action_name]} \n"
                )

        return res

    def teardown(self, stage):
        self._action_count = {}
        self._action_first_occurrence = {}
        super().teardown(stage=stage)
trainer = Trainer(profiler=ActionCountProfiler())
trainer.fit(...)

Profile custom actions of interest

To profile a specific action of interest, reference a profiler in the LightningModule.

from pytorch_lightning.profilers import SimpleProfiler, PassThroughProfiler


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

To profile in any part of your code, use the self.profiler.profile() function

class MyModel(LightningModule):
    def custom_processing_step(self, data):
        with self.profiler.profile("my_custom_action"):
            ...
        return data

Here’s the full code:

from pytorch_lightning.profilers import SimpleProfiler, PassThroughProfiler


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

    def custom_processing_step(self, data):
        with self.profiler.profile("my_custom_action"):
            ...
        return data


profiler = SimpleProfiler()
model = MyModel(profiler)
trainer = Trainer(profiler=profiler, max_epochs=1)

© Copyright Copyright (c) 2018-2022, Lightning AI et al... Revision 7c8c996f.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
Versions
latest
stable
1.7.1
1.7.0
1.6.5
1.6.4
1.6.3
1.6.2
1.6.1
1.6.0
1.5.10
1.5.9
1.5.8
1.5.7
1.5.6
1.5.5
1.5.4
1.5.3
1.5.2
1.5.1
1.5.0
1.4.9
1.4.8
1.4.7
1.4.6
1.4.5
1.4.4
1.4.3
1.4.2
1.4.1
1.4.0
1.3.8
1.3.7
1.3.6
1.3.5
1.3.4
1.3.3
1.3.2
1.3.1
1.3.0
1.2.10
1.2.8
1.2.7
1.2.6
1.2.5
1.2.4
1.2.3
1.2.2
1.2.1
1.2.0
1.1.8
1.1.7
1.1.6
1.1.5
1.1.4
1.1.3
1.1.2
1.1.1
1.1.0
1.0.8
1.0.7
1.0.6
1.0.5
1.0.4
1.0.3
1.0.2
1.0.1
1.0.0
0.10.0
0.9.0
0.8.5
0.8.4
0.8.3
0.8.2
0.8.1
0.8.0
0.7.6
0.7.5
0.7.4
0.7.3
0.7.2
0.7.1
0.7.0
0.6.0
0.5.3
0.4.9
future-structure
Downloads
html
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.