Integrate Pytorch Profiler in

🚀 Feature Request

Let’s add support profiling of functions using the Pytorch profiler.

Motivation

Help understand if there are bottlenecks in code

Proposal

Possible Use Case

import torch
from torch.utils.data import DataLoader, TensorDataset
from catalyst import dl

# data
num_samples, num_features = int(1e4), int(1e1)
X, y = torch.rand(num_samples, num_features), torch.rand(num_samples)
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=32, num_workers=1)
loaders = {"train": loader, "valid": loader}

# model, criterion, optimizer, scheduler
model = ....  # 
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [3, 6])

# model training
runner = dl.SupervisedRunner()
profiler=torch.profiler.profile(
    schedule=torch.profiler.schedule(
        wait=2,
        warmup=2,
        active=6,
        repeat=1),
    on_trace_ready=tensorboard_trace_handler,
    with_trace=True
runner.train(
    model=model,
    criterion=criterion,
    optimizer=optimizer,
    scheduler=scheduler,
    loaders=loaders,
    logdir="./logdir",
    valid_loader="valid",
    valid_metric="loss",
    minimize_valid_metric=True,
    num_epochs=8,
    profiler=profiler
    verbose=True,
)

Alternatives

We could also integrate the profiler into other parts of IRunner.

Additional context

I don’t know whether older versions of Pytorch (<1.8.1) allow you to pass profiler as the context manager to be used.

Checklist

• feature proposal description
• motivation
• extra proposal context / proposal alternatives review

FAQ

Please review the FAQ before submitting an issue:

• I have read the documentation and FAQ
• I have reviewed the minimal examples section
• I have checked the changelog for main framework updates
• I have read the contribution guide
• I have joined Catalyst slack (#__questions channel) for issue discussion

[ssktotoro]

That’s the easy version to play with. The more annoying bit of what I envision is passing the profiler rather than a flag in order to have it be fully configurable. This means that we validate the profiler after its passed, but that’s not too bad either.

[Scitator]

dear @ssktotoro, thanks for the issue Am I correct, that we could do something like

class IRunner:
    def run(self, *args, **kwargs):
        # here is the trick
        if self.profile:
           with torch.profiler.profile(
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
        
            # In this example with wait=1, warmup=1, active=2,
            # profiler will skip the first step/iteration,
            # start warming up on the second, record
            # the third and the forth iterations,
            # after which the trace will become available
            # and on_trace_ready (when set) is called;
            # the cycle repeats starting with the next step
        
            schedule=torch.profiler.schedule(
                wait=1,
                warmup=1,
                active=2),
            on_trace_ready=trace_handler
            # on_trace_ready=torch.profiler.tensorboard_trace_handler('./log')
            # used when outputting for tensorboard
            ) as p:
                self.run_experiment()
        else:
            self.run_experiment()

? So, just make a code injection here or. here?

The example is taken from the PyTorch docs https://pytorch.org/docs/master/profiler.html.