Add PyTorch Profiler in Tensorboard loging

Author: ggrzeczkowicz

README.md

@@ -16,7 +16,7 @@ See our [paper](https://arxiv.org/pdf/2010.04642.pdf) at 3DV for an overview of
 
 - CUDA 10 or higher (if you want GPU version)
 - Python 3.7 or higher + headers (python-dev)
-- PyTorch 1.7 or higher (PyTorch >= 1.9 is recommended)
+- PyTorch 1.8.1 or higher (PyTorch >= 1.9 is recommended)
 - A Sparse convolution backend (optional) see [here](https://github.com/nicolas-chaulet/torch-points3d#3d-sparse-convolution-support) for installation instructions
 
 Install with

conf/training/default.yaml

@@ -41,3 +41,4 @@ wandb:
   # parameters for TensorBoard Visualization
 tensorboard:
   log: True
+  pytorch_profiler: True

docker/install_python.sh

@@ -10,13 +10,13 @@ python3 -m pip install -U pip
 pip3 install setuptools>=41.0.0
 if [ $1 == "gpu" ]; then
     echo "Install GPU"
-    pip3 install torch==1.7.0 torchvision==0.8.1
+    pip3 install torch==1.8.1 torchvision==0.8.1
     pip3 install MinkowskiEngine --install-option="--force_cuda" --install-option="--cuda_home=/usr/local/cuda"
     pip3 install git+https://github.com/mit-han-lab/torchsparse.git@v1.4.0 -v
     pip3 install pycuda
 else
     echo "Install CPU"
-    pip3 install torch==1.7.0+cpu torchvision==0.8.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+    pip3 install torch==1.8.1+cpu torchvision==0.8.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
     pip3 install MinkowskiEngine
     pip3 install git+https://github.com/mit-han-lab/torchsparse.git@v1.4.0
 fi

pyproject.toml

@@ -18,7 +18,7 @@ tqdm = "^4.40"
 open3d = "0.12.0"
 torchnet = "^0.0.4"
 tensorboard = "^2.1"
-torch = "^1.7.0"
+torch = "^1.8.1"
 torch-scatter = "^2.0.0"
 torch-sparse = "^0.6.10"
 torch-cluster = "^1.5.6"

requirements.txt

@@ -122,7 +122,7 @@ tensorboard-plugin-wit==1.8.0
 terminado==0.10.0
 testpath==0.5.0
 threadpoolctl==2.1.0
-torch==1.7.0
+torch==1.8.1
 torch-cluster==1.5.9
 torch-geometric==1.7.1
 torch-points-kernels==0.7.0

torch_points3d/trainer.py

@@ -26,6 +26,10 @@
 from torch_points3d.utils.wandb_utils import Wandb
 from torch_points3d.visualization import Visualizer
 
+# PyTorch Profiler import
+import torch.profiler
+from contextlib import nullcontext
+
 log = logging.getLogger(__name__)
 
 
@@ -136,27 +140,37 @@ def _initialize_trainer(self):
     def train(self):
         self._is_training = True
 
-        for epoch in range(self._checkpoint.start_epoch, self._cfg.training.epochs):
-            log.info("EPOCH %i / %i", epoch, self._cfg.training.epochs)
+        with (torch.profiler.profile(
+            activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
+            schedule=torch.profiler.schedule(skip_first=5, wait=0, warmup=1, active=3),
+            on_trace_ready=torch.profiler.tensorboard_trace_handler(self._tracker._tensorboard_dir),
+            record_shapes=True,
+            profile_memory=True,
+            with_stack=True
+        ) if self.pytorch_profiler_log else nullcontext()) as prof:
 
-            self._train_epoch(epoch)
+            for epoch in range(self._checkpoint.start_epoch, self._cfg.training.epochs):
+                log.info("EPOCH %i / %i", epoch, self._cfg.training.epochs)
 
-            if self.profiling:
-                return 0
+                with (torch.profiler.record_function('train_epoch') if self.pytorch_profiler_log else nullcontext()):
+                    self._train_epoch(epoch, prof)
 
-            if epoch % self.eval_frequency != 0:
-                continue
+                if self.profiling:
+                    return 0
 
-            if self._dataset.has_val_loader:
-                self._test_epoch(epoch, "val")
+                if epoch % self.eval_frequency != 0:
+                    continue
 
-            if self._dataset.has_test_loaders:
-                self._test_epoch(epoch, "test")
+                if self._dataset.has_val_loader:
+                    self._test_epoch(epoch, "val")
 
-        # Single test evaluation in resume case
-        if self._checkpoint.start_epoch > self._cfg.training.epochs:
-            if self._dataset.has_test_loaders:
-                self._test_epoch(epoch, "test")
+                if self._dataset.has_test_loaders:
+                    self._test_epoch(epoch, "test")
+
+            # Single test evaluation in resume case
+            if self._checkpoint.start_epoch > self._cfg.training.epochs:
+                if self._dataset.has_test_loaders:
+                    self._test_epoch(epoch, "test")
 
     def eval(self, stage_name=""):
         self._is_training = False
@@ -180,7 +194,7 @@ def _finalize_epoch(self, epoch):
             if self._tracker._stage == "train":
                 log.info("Learning rate = %f" % self._model.learning_rate)
 
-    def _train_epoch(self, epoch: int):
+    def _train_epoch(self, epoch: int, prof: torch.profiler.profile):
 
         self._model.train()
         self._tracker.reset("train")
@@ -210,6 +224,9 @@ def _train_epoch(self, epoch: int):
 
                 iter_data_time = time.time()
 
+                if self.pytorch_profiler_log:
+                    prof.step()
+
                 if self.early_break:
                     break
 
@@ -314,6 +331,13 @@ def tensorboard_log(self):
         else:
             return False
 
+    @property
+    def pytorch_profiler_log(self):
+        if self.tensorboard_log:
+            return getattr(self._cfg.training.tensorboard, "pytorch_profiler", False)
+        else:
+            return False
+
     @property
     def tracker_options(self):
         return self._cfg.get("tracker_options", {})