Refactor PyTorch profiler 4/5 (#6349)

Co-authored-by: thomas chaton <[email protected]>

Author: carmocca

CHANGELOG.md

@@ -52,6 +52,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `AbstractProfiler` interface ([#6621](https://github.com/PyTorchLightning/pytorch-lightning/pull/6621))
 
 
+- Added support for including module names for forward in the autograd trace of `PyTorchProfiler` ([#6349](https://github.com/PyTorchLightning/pytorch-lightning/pull/6349))
+
+
 - Added `outputs` parameter to callback's `on_validation_epoch_end` & `on_test_epoch_end` hooks ([#6120](https://github.com/PyTorchLightning/pytorch-lightning/pull/6120))
 
 
@@ -72,6 +75,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Changed profilers to save separate report files per state and rank ([#6621](https://github.com/PyTorchLightning/pytorch-lightning/pull/6621))
 
 
+- Changed `PyTorchProfiler` to use `torch.autograd.profiler.record_function` to record functions ([#6349](https://github.com/PyTorchLightning/pytorch-lightning/pull/6349))
+
+
 ### Deprecated
 
 - `period` has been deprecated in favor of `every_n_val_epochs` in the `ModelCheckpoint` callback ([#6146](https://github.com/PyTorchLightning/pytorch-lightning/pull/6146))
@@ -83,6 +89,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Deprecated `Profiler(output_filename)` in favor of `dirpath` and `filename` ([#6621](https://github.com/PyTorchLightning/pytorch-lightning/pull/6621))
 
 
+- Deprecated `PytorchProfiler(profiled_functions)` in favor of `record_functions` ([#6349](https://github.com/PyTorchLightning/pytorch-lightning/pull/6349))
+
+
 - Deprecated metrics in favor of `torchmetrics` ([#6505](https://github.com/PyTorchLightning/pytorch-lightning/pull/6505),
 
     [#6530](https://github.com/PyTorchLightning/pytorch-lightning/pull/6530),

pytorch_lightning/profiler/profilers.py

@@ -126,16 +126,15 @@ def _prepare_filename(self) -> str:
             filename += f"{self._stage}-"
         filename += str(self.filename)
         if self._local_rank is not None:
-            filename += f"-{self.local_rank}"
+            filename += f"-{self._local_rank}"
         filename += ".txt"
         return filename
 
     def _prepare_streams(self) -> None:
         if self._write_stream is not None:
             return
         if self.filename:
-            dirpath = self.dirpath or self._log_dir
-            filepath = os.path.join(dirpath, self._prepare_filename())
+            filepath = os.path.join(self.dirpath, self._prepare_filename())
             fs = get_filesystem(filepath)
             file = fs.open(filepath, "a")
             self._output_file = file
@@ -175,8 +174,7 @@ def setup(
         self._stage = stage
         self._local_rank = local_rank
         self._log_dir = log_dir
-        if self.dirpath is None:
-            self.dirpath = self._log_dir
+        self.dirpath = self.dirpath or log_dir
 
     def teardown(self, stage: Optional[str] = None) -> None:
         """
@@ -202,8 +200,8 @@ def summary(self) -> str:
         raise NotImplementedError
 
     @property
-    def local_rank(self):
-        return '0' if self._local_rank is None else self._local_rank
+    def local_rank(self) -> int:
+        return 0 if self._local_rank is None else self._local_rank
 
 
 class PassThroughProfiler(BaseProfiler):