Merge pull request #13 from awslabs/pre_commit_fix

running pre-commit

Author: Vikas-kum

smdebug/profiler/python_profiler.py

@@ -157,6 +157,7 @@ def _dump_stats(self, stats_file_path):
         get_logger("smdebug-profiler").info(f"Dumping cProfile stats to {stats_file_path}.")
         pstats.Stats(self._profiler).dump_stats(stats_file_path)
 
+
 class PyinstrumentPythonProfiler(PythonProfiler):
     """Higher level class to oversee profiling specific to Pyinstrument, a third party Python profiler.
     """

smdebug/pytorch/hook.py

@@ -1,11 +1,11 @@
 # Standard Library
 import os
 import time
-from packaging import version
 
 # Third Party
 import torch
 import torch.distributed as dist
+from packaging import version
 
 # First Party
 from smdebug.core.collection import DEFAULT_PYTORCH_COLLECTIONS, CollectionKeys
@@ -258,7 +258,6 @@ def _collect_torch_profiling_data_if_profiler_enabled(self):
                     cpu_thread_start_time=event.cpu_interval.start + self.start_profiler_time_us,
                 )
 
-
     # This hook is invoked by trainer prior to running the forward pass.
     def forward_pre_hook(self, module, inputs):
         # Disable pre-step 0 python profiling if profiling is enabled and if this is step 0.
@@ -289,9 +288,13 @@ def forward_pre_hook(self, module, inputs):
 
             if not self.autograd_profiler_enabled:
                 if version.parse(torch.__version__) <= version.parse("1.5.1"):
-                    torch.autograd._enable_profiler(torch.autograd.ProfilerConfig(self.profiler, False))
-                elif  version.parse(torch.__version__) >= version.parse("1.6"):
-                    torch.autograd._enable_profiler(torch.autograd.ProfilerConfig(self.profiler, False, False))
+                    torch.autograd._enable_profiler(
+                        torch.autograd.ProfilerConfig(self.profiler, False)
+                    )
+                elif version.parse(torch.__version__) >= version.parse("1.6"):
+                    torch.autograd._enable_profiler(
+                        torch.autograd.ProfilerConfig(self.profiler, False, False)
+                    )
                 self.start_profiler_time_us = time.time() * CONVERT_TO_MICROSECS
                 self.autograd_profiler_enabled = True
 
@@ -461,10 +464,11 @@ def _closure_for_registering_backward_hook(self, module):
     def count_parameters(self, model):
         total_params = 0
         for name, parameter in model.named_parameters():
-            if not parameter.requires_grad: continue
+            if not parameter.requires_grad:
+                continue
             param = parameter.numel()
             self.logger.info(f"name:{name} count_params:{param}")
-            total_params+=param
+            total_params += param
         self.logger.info(f"Total Trainable Params: {total_params}")
         return total_params
 
@@ -513,7 +517,6 @@ def register_module(self, module):
         self.has_registered_module = True
         self.count_parameters(module)
 
-
     def register_loss(self, loss_module):
         """Register something like `criterion = nn.CrossEntropyLoss()`."""
         # Typechecking

tests/profiler/pytorch/test_pytorch_profiler.py

@@ -81,4 +81,4 @@ def test_pytorch_profiler(pytorch_profiler_config_parser, out_dir):
     lt.refresh_event_file_list()
     events = lt.get_events(0, time.time() * 1000000)
     print(f"Number of events {len(events)}")
-    assert len(events) == 379
+    assert len(events) == 447

tests/profiler/pytorch/test_pytorch_profiler_rnn.py

@@ -72,5 +72,5 @@ def test_pytorch_profiler_rnn(pytorch_profiler_config_parser, out_dir):
     lt.refresh_event_file_list()
     events = lt.get_events(0, time.time() * 1000000)
     print(f"Number of events {len(events)}")
-    assert len(events) == 59
+    assert len(events) == 61
     shutil.rmtree(out_dir, ignore_errors=True)