How to change the schedule of the built-in PyTorch Profiler?

[dpaleka]

When using PyTorch Profiler in plain PyTorch, one can change the profiling schedule, see e.g. the arguments in the first snippet here:

with torch.profiler.profile(
    schedule=torch.profiler.schedule(
        wait=2,
        warmup=2,
        active=6,
        repeat=1),
    on_trace_ready=tensorboard_trace_handler,
    with_stack=True
) as profiler:

How to do this in PyTorch Lightning? There is nothing about the profiler schedule in the docs.

[akihironitta]

Hi @dpaleka! Haven't tried it myself yet, but PyTorchProfiler has **profiler_kwargs in the constructor, and you should be able to pass in the native PyTorch profiler arguments there:

import torch
from pytorch_lightning import Trainer
from pytorch_lightning.profiler import PyTorchProfiler

profiler = PyTorchProfiler(
    schedule=torch.profiler.schedule(
        wait=2,
        warmup=2,
        active=6,
        repeat=1),
    on_trace_ready=tensorboard_trace_handler,
    with_stack=True,
)
trainer = Trainer(..., profiler=profiler)

Related doc pages:

• https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.profiler.PyTorchProfiler.html#pytorch_lightning.profiler.PyTorchProfiler
• https://pytorch-lightning.readthedocs.io/en/stable/advanced/profiler.html#pytorch-profiler

[dpaleka]

Thank you, but unfortunately this does not work. Let's run the following code:

import torch
from pytorch_lightning import Trainer
from pytorch_lightning.profiler import PyTorchProfiler

profiler = PyTorchProfiler(
    schedule=torch.profiler.schedule(
        wait=2,
        warmup=2,
        active=6,
        repeat=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler,
    with_stack=True,
)
trainer = Trainer(..., profiler=profiler)

In my codebase, I get warnings such as:
[W kineto_shim.cpp:246] Profiler is not initialized: skipping profiling metadata
UserWarning: Incorrect schedule: WARMUP followed by NONE
And in the end the profile does not show up on Tensorboard.

If I replace the above code with the following, everything works correctly and shows up on Tensorboard, however I can't set the schedule.

import torch
from pytorch_lightning import Trainer

trainer = Trainer(..., profiler='pytorch)

Does anyone have an example which sets the schedule parameter correctly?

[dpaleka]

Is there any schedule-related test for https://pytorch-lightning.readthedocs.io/en/stable/_modules/pytorch_lightning/profiler/pytorch.html#PyTorchProfiler? I can't make this work properly.

The default schedule seems to be wait=1, warmup=1, active=3. if a schedule parameter is provided to PyTorchProfiler.__init__, it is replaced with a ScheduleWrapper which does something I don't understand.

Can we transfer this discussion to an issue? The schedule parameter should at least be documented.

[matanninio]

It's a bug. See the issue for details ( #14063 )
you can hack it by commenting out

            if self.is_training:
                self._training_step_reached_end = True
            elif self.is_validating:
                self._validation_step_reached_end = True
            elif self.is_testing:
                self._test_step_reached_end = True
            elif self.is_predicting:
                self._predict_step_reached_end = True

in ScheduleWrapper.__call__

[andife]

Are there any updates to that topic?