Merge: [nnUNet/PyT] Remove profiling

Author: nv-kkudrynski

PyTorch/Segmentation/nnUNet/main.py

@@ -15,7 +15,6 @@
 import ctypes
 import os
 
-import nvidia_dlprof_pytorch_nvtx
 import torch
 from pytorch_lightning import Trainer, seed_everything
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, early_stopping
@@ -30,10 +29,6 @@
 if __name__ == "__main__":
     args = get_main_args()
 
-    if args.profile:
-        nvidia_dlprof_pytorch_nvtx.init()
-        print("Profiling enabled")
-
     if args.affinity != "disabled":
         set_affinity(int(os.getenv("LOCAL_RANK", "0")), args.gpus, mode=args.affinity)
 
@@ -67,7 +62,6 @@
                 mode=args.exec_mode,
                 warmup=args.warmup,
                 dim=args.dim,
-                profile=args.profile,
             )
         ]
     elif args.exec_mode == "train":
@@ -108,11 +102,7 @@
 
     if args.benchmark:
         if args.exec_mode == "train":
-            if args.profile:
-                with torch.autograd.profiler.emit_nvtx():
-                    trainer.fit(model, train_dataloader=data_module.train_dataloader())
-            else:
-                trainer.fit(model, train_dataloader=data_module.train_dataloader())
+            trainer.fit(model, train_dataloader=data_module.train_dataloader())
         else:
             # warmup
             trainer.test(model, test_dataloaders=data_module.test_dataloader())

PyTorch/Segmentation/nnUNet/requirements.txt

@@ -5,5 +5,4 @@ pytorch-lightning==1.3.8
 scikit-learn==1.0
 scikit-image==0.18.3
 pynvml==11.0.0
-numpy==1.21.3
-nvidia_dlprof_pytorch_nvtx==1.7.0
+numpy==1.21.3

PyTorch/Segmentation/nnUNet/scripts/benchmark.py

@@ -29,7 +29,6 @@
 parser.add_argument("--warmup", type=int, default=50, help="Warmup iterations before collecting statistics")
 parser.add_argument("--results", type=str, default="/results", help="Path to results directory")
 parser.add_argument("--logname", type=str, default="perf.json", help="Name of dlloger output")
-parser.add_argument("--profile", action="store_true", help="Enable dlprof profiling")
 
 if __name__ == "__main__":
     args = parser.parse_args()
@@ -45,7 +44,6 @@
     cmd += f"--test_batches {args.test_batches} "
     cmd += f"--warmup {args.warmup} "
     cmd += "--amp " if args.amp else ""
-    cmd += "--profile " if args.profile else ""
     if args.mode == "train":
         cmd += f"--batch_size {args.batch_size} "
     else:

PyTorch/Segmentation/nnUNet/utils/logger.py

@@ -17,7 +17,6 @@
 
 import dllogger as logger
 import numpy as np
-import torch.cuda.profiler as profiler
 from dllogger import JSONStreamBackend, StdOutBackend, Verbosity
 from pytorch_lightning import Callback
 
@@ -49,20 +48,17 @@ def flush(self):
 
 
 class LoggingCallback(Callback):
-    def __init__(self, log_dir, filnename, global_batch_size, mode, warmup, dim, profile):
+    def __init__(self, log_dir, filnename, global_batch_size, mode, warmup, dim):
         self.dllogger = DLLogger(log_dir, filnename)
         self.warmup_steps = warmup
         self.global_batch_size = global_batch_size
         self.step = 0
         self.dim = dim
         self.mode = mode
-        self.profile = profile
         self.timestamps = []
 
     def do_step(self):
         self.step += 1
-        if self.profile and self.step == self.warmup_steps:
-            profiler.start()
         if self.step > self.warmup_steps:
             self.timestamps.append(time.time())
 
@@ -96,8 +92,6 @@ def _log(self):
         self.dllogger.flush()
 
     def on_train_end(self, trainer, pl_module):
-        if self.profile:
-            profiler.stop()
         self._log()
 
     def on_test_end(self, trainer, pl_module):