Runner profiler update

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.nn.modules 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()

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.callback import Callback
 from catalyst.core.logger import ILogger
@@ -191,6 +192,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
@@ -303,6 +305,22 @@ 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"),
+                loader_key="train",
+                profiler_kwargs=dict(
+                    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(
@@ -344,6 +362,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,
@@ -385,6 +404,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``.
@@ -510,6 +531,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()