Runner profiler update

.github/workflows/dl_cpu.yml

@@ -69,6 +69,7 @@ jobs:
           python -c "req = open('./requirements/requirements-ml.txt').read().replace('>', '=') ; open('./requirements/requirements-ml.txt', 'w').write(req)"
           python -c "req = open('./requirements/requirements-neptune.txt').read().replace('>', '=') ; open('./requirements/requirements-neptune.txt', 'w').write(req)"
           python -c "req = open('./requirements/requirements-albu.txt').read().replace('>', '=') ; open('./requirements/requirements-albu.txt', 'w').write(req)"
+          python -c "req = open('./requirements/requirements-profiler.txt').read().replace('>', '=') ; open('./requirements/requirements-profiler.txt', 'w').write(req)"
 
       # https://github.com/actions/cache/blob/master/examples.md
       # Note: This uses an internal pip API and may not always work
@@ -97,7 +98,7 @@ jobs:
       - name: install dependencies
         run: |
           # python -m pip install --upgrade --user pip
-          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-neptune.txt -r ./requirements/requirements-albu.txt
+          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-neptune.txt -r ./requirements/requirements-albu.txt -r ./requirements/requirements-profiler.txt
           python --version
           pip --version
           pip list

catalyst/callbacks/profiler.py

@@ -51,7 +51,6 @@ class ProfilerCallback(Callback):
             from catalyst import dl
             from catalyst.data import ToTensor
             from catalyst.contrib.datasets import MNIST
-            from catalyst.contrib.layers import Flatten
 
             loaders = {
                 "train": DataLoader(
@@ -64,7 +63,7 @@ class ProfilerCallback(Callback):
                 ),
             }
 
-            model = nn.Sequential(Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10))
+            model = nn.Sequential(nn.Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10))
             criterion = nn.CrossEntropyLoss()
             optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
             runner = dl.SupervisedRunner()
@@ -122,16 +121,6 @@ def __init__(
         self.profiler = None
         self.stats = None
 
-    def on_experiment_start(self, runner: IRunner) -> None:
-        """
-        On batch end action
-
-        Args:
-            runner: current runner
-        """
-        if self.loader_key is None:
-            self.loader_key = runner.loader_key  # use first loader for profile
-
     def _should_use_profiler(self, loader_key: str, epoch: int):
         if self.loader_key == loader_key and self.epoch == epoch:
             if self.num_batches is not None:
@@ -186,6 +175,9 @@ def on_loader_start(self, runner: IRunner) -> None:
         Args:
             runner: current runner
         """
+        if self.loader_key is None:
+            self.loader_key = runner.loader_key  # use first loader for profile
+
         self._enter_profiler(runner)
 
     def on_loader_end(self, runner: IRunner) -> None:

catalyst/runners/runner.py

@@ -12,6 +12,7 @@
 from catalyst.callbacks.criterion import CriterionCallback, ICriterionCallback
 from catalyst.callbacks.misc import CheckRunCallback, TimerCallback, TqdmCallback
 from catalyst.callbacks.optimizer import IOptimizerCallback, OptimizerCallback
+from catalyst.callbacks.profiler import ProfilerCallback
 from catalyst.callbacks.scheduler import ISchedulerCallback, SchedulerCallback
 from catalyst.core._misc import callback_isinstance, sort_callbacks_by_order
 from catalyst.core.callback import Callback
@@ -192,6 +193,7 @@ def __init__(self, *args, **kwargs):
         self._timeit = False
         self._check = False
         self._overfit = False
+        self._profile = False
         self._load_best_on_end = False
 
     @property
@@ -304,6 +306,21 @@ def get_callbacks(self, stage: str) -> "OrderedDict[str, Callback]":
             callbacks["_check"] = CheckRunCallback()
         if self._overfit and not is_callback_exists(BatchOverfitCallback):
             callbacks["_overfit"] = BatchOverfitCallback()
+        if self._profile and not is_callback_exists(ProfilerCallback):
+            callbacks["_profile"] = ProfilerCallback(
+                tensorboard_path=os.path.join(self._logdir, "tb_profile"),
+                profiler_kwargs={
+                    "activities": [
+                        torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.CUDA,
+                    ],
+                    "on_trace_ready": torch.profiler.tensorboard_trace_handler(
+                        os.path.join(self._logdir, "tb_profile")
+                    ),
+                    "with_stack": True,
+                    "with_flops": True,
+                },
+            )
 
         if self._logdir is not None and not is_callback_exists(ICheckpointCallback):
             callbacks["_checkpoint"] = CheckpointCallback(
@@ -348,6 +365,7 @@ def train(
         timeit: bool = False,
         check: bool = False,
         overfit: bool = False,
+        profile: bool = False,
         load_best_on_end: bool = False,
         # engine extra params,
         fp16: bool = False,
@@ -390,6 +408,8 @@ def train(
             overfit: if True, then takes only one batch per loader
                 for model overfitting, for advance usage please check
                 ``BatchOverfitCallback``
+            profile: if True, then uses ProfilerCallback, for advance usage please check
+                ``ProfilerCallback``
             load_best_on_end: if True, Runner will load
                 best checkpoint state (model, optimizer, etc)
                 according to validation metrics. Requires specified ``logdir``.
@@ -516,6 +536,7 @@ def on_loader_end(self, runner):
         self._timeit = timeit
         self._check = check
         self._overfit = overfit
+        self._profile = profile
         self._load_best_on_end = load_best_on_end
         # run
         self.run()

docs/faq/settings.rst

@@ -24,6 +24,7 @@ Nevertheless, there are much more availabe:
     pip install catalyst[onnx-gpu] # + onnx, onnxruntime-gpu
     pip install catalyst[onnx] # + onnx, onnxruntime
     pip install catalyst[optuna] # + optuna
+    pip install catalyst[profiler] # + profiler
     pip install catalyst[wandb] # + wandb
     pip install catalyst[all] # + catalyst[cv], catalyst[ml], catalyst[optuna]
 

requirements/requirements-profiler.txt

@@ -0,0 +1 @@
+torch_tb_profiler

setup.py

@@ -54,6 +54,7 @@ def load_version():
     "onnx-gpu": load_requirements("requirements/requirements-onnx-gpu.txt"),
     "onnx": load_requirements("requirements/requirements-onnx.txt"),
     "optuna": load_requirements("requirements/requirements-optuna.txt"),
+    "profiler": load_requirements("requirements/requirements-profiler.txt"),
     "wandb": load_requirements("requirements/requirements-wandb.txt"),
     # "xla": load_requirements("requirements/requirements-xla.txt"),
 }

tests/pipelines/test_runner.py

@@ -0,0 +1,54 @@
+# flake8: noqa
+
+import os
+
+from pytest import mark
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+from catalyst import dl
+from catalyst.contrib.data import ImageToTensor
+from catalyst.contrib.datasets import MNIST
+
+
+def _is_profile_available():
+    try:
+        from torch import profiler  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+def train_experiment():
+    loaders = {
+        "train": DataLoader(
+            MNIST(os.getcwd(), train=False, download=True, transform=ImageToTensor()),
+            batch_size=32,
+        ),
+        "valid": DataLoader(
+            MNIST(os.getcwd(), train=False, download=True, transform=ImageToTensor()),
+            batch_size=32,
+        ),
+    }
+    model = nn.Sequential(nn.Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10))
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
+    runner = dl.SupervisedRunner()
+    runner.train(
+        model=model,
+        loaders=loaders,
+        criterion=criterion,
+        optimizer=optimizer,
+        num_epochs=5,
+        logdir="./logs",
+        profile=True,
+    )
+
+
+@mark.skipif(not _is_profile_available(), reason="Torch profiler is not available")
+def test_profiler():
+    train_experiment()
+    assert os.path.isdir("./logs/tb_profile") and not len(os.listdir("./logs/tb_profile")) == 0