Repository: Lightning-AI/pytorch-lightning
Pull request: new feature for profiling training runs
# Before submitting

- [x] Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md)?
- [x] Did you make sure to update the docs?   
- [x] Did you write any new necessary tests? 


## What does this PR do?
Fixes #753 

## PR review    
Anyone in the community is free to review the PR once the tests have passed.     
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.



Problem statement:
Feature: Profiling for a training run
## 🚀 Feature
It'd be nice if the PyTorch Lightning Trainer had a way for profiling a training run so that I could easily identify where bottlenecks are occurring. 

### Motivation

I have been developing a model and had been using a small toy dataset to verify that everything is working properly. When I started a training run on a larger dataset, the estimated epoch duration was much longer than I anticipated. I have many ideas about where the bottleneck might be occurring, but I don't have an easy way to validate with data.

### Pitch

We should include a flag when constructing the Trainer which signifies that we would like to profile this training run. The default will be False since the profiling output may be quite verbose and unnecessary in most cases.

```
trainer = Trainer(args, profile=True)
trainer.fit(model)
```

I was initially thinking that we could write a simple decorator and apply this on all the actions (ie functions) which we wish to profile. The benefit of writing a decorator is that the scope of an action is clearly defined (so we know when to start and stop the timing). 

Something along the lines of (non functional example):

```
from functools import wraps
import time


def profile(func, enabled=False):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        output = func(*args, **kwargs)
        end = time.time()
        print(f'{func.__name__} took {end-start}s')
        return output
    return wrapper


class Model:
    def __init__(self):
        pass

    @profile(enabled=trainer.profile)
    def training_step(self, batch, batch_idx):
        pass

    @profile(enabled=trainer.profile)
    def validation_step(self, batch, batch_idx):
        pass
```

Some issues with this approach:
- We want to define whether or not a profile is enabled when constructing the trainer, so it's not easy to access this when decorating functions within the model. 
- Users are expected to overwrite these methods, so we won't actually be able to reliably decorate them anyways. 

@williamFalcon suggested (and please correct me if I misstate something) that we could construct a profiler object if `profile=True` in the Trainer initialization. We could then inject `self.profiler.profile()` in the training loop code for every major event (get train data, get val data, training step, backward, optimizer step, val step, etc). 

As I understand it, the `profiler` would be something along the lines of:

```
class Profiler:
    def __init__(self):
        self.last_recorded_time = None
        self.last_recorded_action = None

    def profile(self, action=''):
        current_time = time.time()
        if self.last_recorded_time is not None and self.last_recorded_action is not None:
            print(f'{self.last_recorded_action} took {current_time - self.last_recorded_time}s'
        self.last_recorded_time = current_time
        self.last_recorded_action = action
```

My concern with this approach (as I currently understand it) is that we define the "scope" of an action as simply waiting until the next invocation of `profile`. Especially given the fact that functions are calling other functions during training (eg. `run_training_epoch` -> `run_training_batch`) it can be very hard to keep track of the scope of an action. 

### Alternatives

One way we can address this is by invoking the profile within a context handler, but this might get messy by effectively introducing another level of indentation for most of the code.

```
with profile():
    # do action here
```

Hints from issue discussion:
Thanks to Github's suggestion, I found this similar issue #251 
Great suggestion. I still prefer the logger approach instead of decorator.

### API
I imagine it like:

```
Trainer(profile=True|False|Profiler)
```

- If True, we use default Profiler
```
self.profiler = Profiler(...)
```
- If False, no profiler

- Otherwise use the give profiler

### Trainer
In  the trainer, we'd call it at key places

```
tng_data = get_train_data()
self.profile('train data')

for epoch in epochs:
  for batch in dataset:
      training_step()
      self.profile('training step')

      backward()
      self.profile('backward')

      #...  etc
```

### Usability
We also allow users to call it whenever they want

```
def training_step(...):
    x, y = batch
    x = self.forward(x)
    self.profile('fwd')

    x = self.second_nn(x)
    self.profile('projection a')
    
    x = self.loss(x)
    self.profile('loss')
```
for the record, more discussion is here: 
https://pytorch-lightning.slack.com/archives/CRBLFHY79/p1579997800007700
What if we did the logger approach but made the scope more explicit?

Eg.
```
self.profile.start('fetch train data')
tng_data = get_train_data()
self.profile.stop('fetch train data')

for epoch in epochs:
  for batch in dataset:
      self.profile.start('training step')
      training_step()
      self.profile.stop('training step')

      self.profile.start('backward')
      backward()
      self.profile.stop('backward')

      #...  etc
```

And then we could expose an alternative interface using context handlers if the user prefers. 

```
with self.profile('fetch train data'):
  tng_data = get_train_data()

for epoch in epochs:
  for batch in dataset:
      with self.profile('training step'):
          training_step()

      with self.profile('backward')
          backward()

      #...  etc
```

Internally, we can store the values as a `defaultdict(list)` and append each (stop-start) duration, so that at the end of the run we can report mean times for each action as well as standard deviations. 
I'm also curious what design @jeffling decided on for their profiling tool
@jeremyjordan I like the context handler suggestion. That way, I can decide on my won contexts.

Want to take a stab at implementing this?
@jeremyjordan 

We use several, but the most useful one was really just using a lot of @timeits. We have a small edit that forces GPU sync so we can get GPU time. Unfortunately, still clunky to use but works. We expose both a decorator and logger api, and the logger one definitely is more popular with the team. 

Took an initial stab at implementing this yesterday evening (see branch [here](https://github.com/jeremyjordan/pytorch-lightning/blob/feature/profiling/pytorch_lightning/utilities/profiler.py)).

I wrote a simple profiler which just records durations of actions:

```
>>> p = Profiler()
>>> 
>>> with p.profile("context handler"):
...     time.sleep(5)
...     a = np.random.randn(3000, 2)
...     b = a + 2
...     c = b / 3
... 
>>> p.start('manual')
>>> time.sleep(5)
>>> a = np.random.randn(3000, 2)
>>> b = a + 2
>>> c = b / 3
>>> p.stop('manual')
>>> 
>>> p.describe()
Action                  |  Mean duration (s)    |  Std dev.       
------------------------------------------------------------
context handler         |  5.0022               |  0.0            
manual                  |  5.0036               |  0.0 
```

And a more advanced profiler which leverages Python's cProfiler which recorded a more detailed view:

```
>>> ap = AdvancedProfiler()
>>> with ap.profile("test"):
...     time.sleep(5)
...     a = 5*3
...     b = a + 2
...     c = b / 3
... 
>>> ap.start('ok')
>>> a = 3*2
>>> ap.stop('ok')
>>> ap.describe()
Profile stats for: test
         7 function calls in 5.000 seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    5.000    5.000    5.000    5.000 {built-in method time.sleep}
        1    0.000    0.000    0.000    0.000 contextlib.py:116(__exit__)
        1    0.000    0.000    0.000    0.000 {built-in method builtins.next}
        1    0.000    0.000    0.000    0.000 profiler.py:37(profile)
        1    0.000    0.000    0.000    0.000 profiler.py:109(stop)
        1    0.000    0.000    0.000    0.000 {method 'get' of 'dict' objects}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}



Profile stats for: ok
         4 function calls in 0.000 seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.000    0.000 profiler.py:109(stop)
        1    0.000    0.000    0.000    0.000 <stdin>:1(<module>)
        1    0.000    0.000    0.000    0.000 {method 'get' of 'dict' objects}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

```

Still need to validate that everything works properly but figured I'd share the branch in case anyone wants to follow along with the progress.

A couple things I might need help on:
- I haven't done much distributed training in PyTorch so I might want to get some review and make sure this will work properly with that.
- I'm not sure I understand the motivation behind a GPU sync to get GPU time
Also, at what point do we want to print the results from the training profile?
@jeremyjordan could pls make it as PR and add "wip" to its name so it is easier to check the chances...

Commit message: option to dump cProfiler results to text file

Patch to review:
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
index 0b15100606f6c..7026dd2aeefcf 100644
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -26,6 +26,8 @@
 from pytorch_lightning.trainer.training_loop import TrainerTrainLoopMixin
 from pytorch_lightning.trainer.training_tricks import TrainerTrainingTricksMixin
 from pytorch_lightning.utilities.debugging import MisconfigurationException
+from pytorch_lightning.utilities.profiler import PassThroughProfiler
+
 
 try:
     from apex import amp
@@ -87,6 +89,7 @@ def __init__(
             num_sanity_val_steps=5,
             truncated_bptt_steps=None,
             resume_from_checkpoint=None,
+            profiler=None
     ):
         r"""
 
@@ -460,6 +463,16 @@ def __init__(
 
                     # resume from a specific checkpoint
                     trainer = Trainer(resume_from_checkpoint='some/path/to/my_checkpoint.ckpt')
+            profiler (utilities.BaseProfiler):  To profile individual steps during training and assist in
+                identifying bottlenecks.
+                Example::
+
+                    # default used by the Trainer
+                    trainer = Trainer(profiler=None)
+
+                    # profile a training run and get a report on completion of the training job
+                    profiler = utilities.Profiler()
+                    trainer = Trainer(profiler=profiler)
 
         .. warning:: Following arguments become deprecated and they will be removed in v0.8.0:
 
@@ -564,6 +577,9 @@ def __init__(
         # configure logger
         self.configure_logger(logger)
 
+        # configure profiler
+        self.profiler = profiler or PassThroughProfiler()
+
         # configure early stop callback
         # creates a default one if none passed in
         self.configure_early_stopping(early_stop_callback)
@@ -884,6 +900,9 @@ def run_pretrain_routine(self, model):
         # CORE TRAINING LOOP
         self.train()
 
+        # summarize profile results
+        self.profiler.describe()
+
     def test(self, model=None):
         r"""
 
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
index e1f90308a76e5..b598ec89a1934 100644
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -211,6 +211,7 @@ def __init__(self):
         self.training_tqdm_dict = None
         self.get_train_dataloader = None
         self.reduce_lr_on_plateau_scheduler = None
+        self.profiler = None
 
     @property
     def max_nb_epochs(self):
@@ -355,12 +356,14 @@ def train(self):
                 stop = should_stop and met_min_epochs
                 if stop:
                     self.main_progress_bar.close()
-                    model.on_train_end()
+                    with self.profiler.profile('on_train_end'):
+                        model.on_train_end()
                     return
 
         self.main_progress_bar.close()
 
-        model.on_train_end()
+        with self.profiler.profile('on_train_end'):
+            model.on_train_end()
 
         if self.logger is not None:
             self.logger.finalize("success")
@@ -369,7 +372,8 @@ def run_training_epoch(self):
         # before epoch hook
         if self.is_function_implemented('on_epoch_start'):
             model = self.get_model()
-            model.on_epoch_start()
+            with self.profiler.profile('on_epoch_start'):
+                model.on_epoch_start()
 
         # run epoch
         for batch_idx, batch in enumerate(self.get_train_dataloader()):
@@ -430,7 +434,8 @@ def run_training_epoch(self):
         # epoch end hook
         if self.is_function_implemented('on_epoch_end'):
             model = self.get_model()
-            model.on_epoch_end()
+            with self.profiler.profile('on_epoch_end'):
+                model.on_epoch_end()
 
     def run_training_batch(self, batch, batch_idx):
         # track grad norms
@@ -448,7 +453,8 @@ def run_training_batch(self, batch, batch_idx):
         # hook
         if self.is_function_implemented('on_batch_start'):
             model_ref = self.get_model()
-            response = model_ref.on_batch_start(batch)
+            with self.profiler.profile('on_batch_start'):
+                response = model_ref.on_batch_start(batch)
 
             if response == -1:
                 return -1, grad_norm_dic, {}
@@ -456,7 +462,8 @@ def run_training_batch(self, batch, batch_idx):
         splits = [batch]
         if self.truncated_bptt_steps is not None:
             model_ref = self.get_model()
-            splits = model_ref.tbptt_split_batch(batch, self.truncated_bptt_steps)
+            with self.profiler.profile('tbptt_split_batch'):
+                splits = model_ref.tbptt_split_batch(batch, self.truncated_bptt_steps)
 
         self.hiddens = None
         for split_idx, split_batch in enumerate(splits):
@@ -476,8 +483,9 @@ def run_training_batch(self, batch, batch_idx):
                 # wrap the forward step in a closure so second order methods work
                 def optimizer_closure():
                     # forward pass
-                    output = self.training_forward(
-                        split_batch, batch_idx, opt_idx, self.hiddens)
+                    with self.profiler.profile('model_forward'):
+                        output = self.training_forward(
+                            split_batch, batch_idx, opt_idx, self.hiddens)
 
                     closure_loss = output[0]
                     progress_bar_metrics = output[1]
@@ -491,7 +499,8 @@ def optimizer_closure():
 
                     # backward pass
                     model_ref = self.get_model()
-                    model_ref.backward(self.use_amp, closure_loss, optimizer, opt_idx)
+                    with self.profiler.profile('model_backward'):
+                        model_ref.backward(self.use_amp, closure_loss, optimizer, opt_idx)
 
                     # track metrics for callbacks
                     all_callback_metrics.append(callback_metrics)
@@ -503,7 +512,8 @@ def optimizer_closure():
                     # insert after step hook
                     if self.is_function_implemented('on_after_backward'):
                         model_ref = self.get_model()
-                        model_ref.on_after_backward()
+                        with self.profiler.profile('on_after_backward'):
+                            model_ref.on_after_backward()
 
                     return closure_loss
 
@@ -533,8 +543,9 @@ def optimizer_closure():
                     # calls .step(), .zero_grad()
                     # override function to modify this behavior
                     model = self.get_model()
-                    model.optimizer_step(self.current_epoch, batch_idx,
-                                         optimizer, opt_idx, optimizer_closure)
+                    with self.profiler.profile('optimizer_step'):
+                        model.optimizer_step(self.current_epoch, batch_idx,
+                                             optimizer, opt_idx, optimizer_closure)
 
                     # calculate running loss for display
                     self.running_loss.append(self.batch_loss_value)
@@ -544,7 +555,8 @@ def optimizer_closure():
         # activate batch end hook
         if self.is_function_implemented('on_batch_end'):
             model = self.get_model()
-            model.on_batch_end()
+            with self.profiler.profile('on_batch_end'):
+                model.on_batch_end()
 
         # update progress bar
         self.main_progress_bar.update(1)
@@ -604,7 +616,8 @@ def training_forward(self, batch, batch_idx, opt_idx, hiddens):
         # allow any mode to define training_end
         if self.is_overriden('training_end'):
             model_ref = self.get_model()
-            output = model_ref.training_end(output)
+            with self.profiler.profile('training_end'):
+                output = model_ref.training_end(output)
 
         # format and reduce outputs accordingly
         output = self.process_output(output, train=True)
diff --git a/pytorch_lightning/utilities/profiler.py b/pytorch_lightning/utilities/profiler.py
new file mode 100644
index 0000000000000..418b1dc8f7ec2
--- /dev/null
+++ b/pytorch_lightning/utilities/profiler.py
@@ -0,0 +1,178 @@
+"""
+Profiling your training run can help you understand if there are any bottlenecks in your code.
+
+PyTorch Lightning supports profiling standard actions in the training loop out of the box, including:
+- on_epoch_start
+- on_epoch_end
+- on_batch_start
+- tbptt_split_batch
+- model_forward
+- model_backward
+- on_after_backward
+- optimizer_step
+- on_batch_end
+- training_end
+- on_training_end
+
+If you only wish to profile the standard actions, you can construct a Profiler object and simply
+pass it into the Trainer.
+
+.. code-block:: python
+    profiler = Profiler()
+    trainer = Trainer(..., profiler=profiler)
+
+You can also reference this profiler to profiler any arbitrary code.
+
+.. code-block:: python
+    with profiler.profile('my_custom_action'):
+        my_custom_action()
+"""
+
+
+from contextlib import contextmanager
+from collections import defaultdict
+import time
+import numpy as np
+import cProfile
+import pstats
+import io
+from abc import ABC, abstractmethod
+
+
+class BaseProfiler(ABC):
+    @abstractmethod
+    def start(self, action_name):
+        """
+        defines how to start recording an action
+        """
+        pass
+
+    @abstractmethod
+    def stop(self, action_name):
+        """
+        defines how to record the duration once an action is complete
+        """
+        pass
+
+    @contextmanager
+    def profile(self, action_name):
+        """
+        yields a context manager to encapsulate the scope of a profiled action
+
+        with self.profile('load training data'):
+            # load training data code
+
+        the profiler will start once you've entered the context and automatically stop
+        once you exit the code block
+        """
+        try:
+            self.start(action_name)
+            yield action_name
+        finally:
+            self.stop(action_name)
+
+    def describe(self):
+        """
+        prints a report after the conclusion of the profiled training run
+        """
+        pass
+
+
+class PassThroughProfiler(BaseProfiler):
+    """
+    this can be used when you don't want to profile your runs
+    """
+
+    def __init__(self):
+        pass
+
+    def start(self, action_name):
+        pass
+
+    def stop(self, action_name):
+        pass
+
+
+class Profiler(BaseProfiler):
+    """
+    this profiler simply records the duration of actions (in seconds) and reports
+    the mean and standard deviation of each action duration over the entire training run
+    """
+
+    def __init__(self):
+        self.current_actions = {}
+        self.recorded_durations = defaultdict(list)
+
+    def start(self, action_name):
+        if action_name in self.current_actions:
+            raise ValueError(
+                f"Attempted to start {action_name} which has already started."
+            )
+        self.current_actions[action_name] = time.monotonic()
+
+    def stop(self, action_name):
+        end_time = time.monotonic()
+        if action_name not in self.current_actions:
+            raise ValueError(
+                f"Attempting to stop recording an action ({action_name}) which was never started."
+            )
+        start_time = self.current_actions.pop(action_name)
+        duration = end_time - start_time
+        self.recorded_durations[action_name].append(duration)
+
+    def describe(self):
+        def print_row(action, mean, std_dev):
+            print(f"{action:<20s}\t|  {mean:<15}\t|  {std_dev:<15}")
+
+        print_row("Action", "Mean duration (s)", "Std dev.")
+        print("-" * 60)
+        for action, durations in self.recorded_durations.items():
+            print_row(action, f"{np.mean(durations):.5}", f"{np.std(durations):.5}")
+
+
+class AdvancedProfiler(BaseProfiler):
+    """
+    this profiler uses Python's cProfiler to record more detailed information about
+    time spent in each function call recorded during a given action
+    """
+
+    def __init__(self, output_filename=None):
+        '''
+        :param output_filename (str): optionally save profile results to file instead of printing
+            to std out when training is finished.
+        '''
+        self.profiled_actions = {}
+        self.output_filename = output_filename
+
+    def start(self, action_name):
+        if action_name not in self.profiled_actions:
+            self.profiled_actions[action_name] = cProfile.Profile()
+        self.profiled_actions[action_name].enable()
+
+    def stop(self, action_name):
+        pr = self.profiled_actions.get(action_name)
+        if pr is None:
+            raise ValueError(
+                f"Attempting to stop recording an action ({action_name}) which was never started."
+            )
+        pr.disable()
+
+    def describe(self, line_count_restriction=1.0):
+        self.recorded_stats = {}
+        for action_name, pr in self.profiled_actions.items():
+            s = io.StringIO()
+            sortby = pstats.SortKey.CUMULATIVE
+            ps = pstats.Stats(pr, stream=s).strip_dirs().sort_stats(sortby)
+            ps.print_stats(line_count_restriction)
+            self.recorded_stats[action_name] = s.getvalue()
+        if self.output_filename is not None:
+            # save to file
+            with open(self.output_filename, 'w') as f:
+                for action, stats in self.recorded_stats.items():
+                    f.write(f"Profile stats for: {action}")
+                    f.write(stats)
+        else:
+            # print to standard out
+            for action, stats in self.recorded_stats.items():
+                print(f"Profile stats for: {action}")
+                print(stats)
diff --git a/tests/test_profiler.py b/tests/test_profiler.py
new file mode 100644
index 0000000000000..76f843d781085
--- /dev/null
+++ b/tests/test_profiler.py
@@ -0,0 +1,24 @@
+from pytorch_lightning.utilities.profiler import Profiler, AdvancedProfiler
+import time
+import numpy as np
+
+
+def test_simple_profiler():
+    p = Profiler()
+
+    with p.profile("a"):
+        time.sleep(3)
+
+    with p.profile("a"):
+        time.sleep(1)
+
+    with p.profile("b"):
+        time.sleep(2)
+
+    with p.profile("c"):
+        time.sleep(1)
+
+    # different environments have different precision when it comes to time.sleep()
+    np.testing.assert_almost_equal(p.recorded_durations["a"], [3, 1], decimal=1)
+    np.testing.assert_almost_equal(p.recorded_durations["b"], [2], decimal=1)
+    np.testing.assert_almost_equal(p.recorded_durations["c"], [1], decimal=1)


Code review comments:
diff_hunk: @@ -0,0 +1,178 @@
+"""
+Profiling your training run can help you understand if there are any bottlenecks in your code.
+
+PyTorch Lightning supports profiling standard actions in the training loop out of the box, including:
+- on_epoch_start
+- on_epoch_end
+- on_batch_start
+- tbptt_split_batch
+- model_forward
+- model_backward
+- on_after_backward
+- optimizer_step
+- on_batch_end
+- training_end
+- on_training_end
+
+If you only wish to profile the standard actions, you can construct a Profiler object and simply
+pass it into the Trainer.
+
+.. code-block:: python
+    profiler = Profiler()
+    trainer = Trainer(..., profiler=profiler)
+
+You can also reference this profiler to profiler any arbitrary code.
+
+.. code-block:: python
+    with profiler.profile('my_custom_action'):
+        my_custom_action()
+"""
+
+
+from contextlib import contextmanager
+from collections import defaultdict
+import time
+import numpy as np
+import cProfile
+import pstats
+import io
+from abc import ABC, abstractmethod
+
+
+class BaseProfiler(ABC):
+    @abstractmethod
+    def start(self, action_name):
+        """
+        defines how to start recording an action
+        """
+        pass
+
+    @abstractmethod
+    def stop(self, action_name):
+        """
+        defines how to record the duration once an action is complete
+        """
+        pass
+
+    @contextmanager
+    def profile(self, action_name):
+        """
+        yields a context manager to encapsulate the scope of a profiled action
+
+        with self.profile('load training data'):
+            # load training data code
+
+        the profiler will start once you've entered the context and automatically stop
+        once you exit the code block
+        """
+        try:
+            self.start(action_name)
+            yield action_name
+        finally:
+            self.stop(action_name)
+
+    def describe(self):
+        """
+        prints a report after the conclusion of the profiled training run
+        """
+        pass
+
+
+class PassThroughProfiler(BaseProfiler):
+    """
+    this can be used when you don't want to profile your runs
+    """
+
+    def __init__(self):
+        pass
+
+    def start(self, action_name):
+        pass
+
+    def stop(self, action_name):
+        pass
+
+
+class Profiler(BaseProfiler):
+    """
+    this profiler simply records the duration of actions (in seconds) and reports
+    the mean and standard deviation of each action duration over the entire training run
+    """
+
+    def __init__(self):
+        self.current_actions = {}
+        self.recorded_durations = defaultdict(list)
+
+    def start(self, action_name):
+        if action_name in self.current_actions:
+            raise ValueError(
+                f"Attempted to start {action_name} which has already started."
+            )
+        self.current_actions[action_name] = time.monotonic()
+
+    def stop(self, action_name):
+        end_time = time.monotonic()
+        if action_name not in self.current_actions:
+            raise ValueError(
+                f"Attempting to stop recording an action ({action_name}) which was never started."
+            )
+        start_time = self.current_actions.pop(action_name)
+        duration = end_time - start_time
+        self.recorded_durations[action_name].append(duration)
+
+    def describe(self):
+        def print_row(action, mean, std_dev):
+            print(f"{action:<20s}\t|  {mean:<15}\t|  {std_dev:<15}")
path: pytorch_lightning/utilities/profiler.py
text: @user1:
use `logging`

diff_hunk: @@ -0,0 +1,178 @@
+"""
+Profiling your training run can help you understand if there are any bottlenecks in your code.
+
+PyTorch Lightning supports profiling standard actions in the training loop out of the box, including:
+- on_epoch_start
+- on_epoch_end
+- on_batch_start
+- tbptt_split_batch
+- model_forward
+- model_backward
+- on_after_backward
+- optimizer_step
+- on_batch_end
+- training_end
+- on_training_end
+
+If you only wish to profile the standard actions, you can construct a Profiler object and simply
+pass it into the Trainer.
+
+.. code-block:: python
+    profiler = Profiler()
+    trainer = Trainer(..., profiler=profiler)
+
+You can also reference this profiler to profiler any arbitrary code.
+
+.. code-block:: python
+    with profiler.profile('my_custom_action'):
+        my_custom_action()
+"""
+
+
+from contextlib import contextmanager
+from collections import defaultdict
+import time
+import numpy as np
+import cProfile
+import pstats
+import io
+from abc import ABC, abstractmethod
+
+
+class BaseProfiler(ABC):
+    @abstractmethod
+    def start(self, action_name):
+        """
+        defines how to start recording an action
+        """
+        pass
+
+    @abstractmethod
+    def stop(self, action_name):
+        """
+        defines how to record the duration once an action is complete
+        """
+        pass
+
+    @contextmanager
+    def profile(self, action_name):
+        """
+        yields a context manager to encapsulate the scope of a profiled action
+
+        with self.profile('load training data'):
+            # load training data code
+
+        the profiler will start once you've entered the context and automatically stop
+        once you exit the code block
+        """
+        try:
+            self.start(action_name)
+            yield action_name
+        finally:
+            self.stop(action_name)
+
+    def describe(self):
+        """
+        prints a report after the conclusion of the profiled training run
+        """
+        pass
+
+
+class PassThroughProfiler(BaseProfiler):
+    """
+    this can be used when you don't want to profile your runs
+    """
+
+    def __init__(self):
+        pass
+
+    def start(self, action_name):
+        pass
+
+    def stop(self, action_name):
+        pass
+
+
+class Profiler(BaseProfiler):
+    """
+    this profiler simply records the duration of actions (in seconds) and reports
+    the mean and standard deviation of each action duration over the entire training run
+    """
+
+    def __init__(self):
+        self.current_actions = {}
+        self.recorded_durations = defaultdict(list)
+
+    def start(self, action_name):
+        if action_name in self.current_actions:
+            raise ValueError(
+                f"Attempted to start {action_name} which has already started."
+            )
+        self.current_actions[action_name] = time.monotonic()
+
+    def stop(self, action_name):
+        end_time = time.monotonic()
+        if action_name not in self.current_actions:
+            raise ValueError(
+                f"Attempting to stop recording an action ({action_name}) which was never started."
+            )
+        start_time = self.current_actions.pop(action_name)
+        duration = end_time - start_time
+        self.recorded_durations[action_name].append(duration)
+
+    def describe(self):
+        def print_row(action, mean, std_dev):
+            print(f"{action:<20s}\t|  {mean:<15}\t|  {std_dev:<15}")
+
+        print_row("Action", "Mean duration (s)", "Std dev.")
+        print("-" * 60)
+        for action, durations in self.recorded_durations.items():
+            print_row(action, f"{np.mean(durations):.5}", f"{np.std(durations):.5}")
+
+
+class AdvancedProfiler(BaseProfiler):
+    """
+    this profiler uses Python's cProfiler to record more detailed information about
+    time spent in each function call recorded during a given action
+    """
+
+    def __init__(self, output_filename=None):
+        '''
path: pytorch_lightning/utilities/profiler.py
text: @user1:
''' >> """