Updated on-gpu benchmaking with cuda event & sync

Author: LukasHedegaard

CHANGELOG.md

@@ -9,20 +9,24 @@ From v1.0.0 and on, the project will adherence strictly to Semantic Versioning.
 
 ## [Unreleased]
 
+## [0.2.1] - 2022-02-11
+### Fixed
+- Updated on-gpu model benchmaking with best-practices on `cuda.Event` and `cuda.synchronize`.
+
 
-## [0.2.0] - 2022-02-11]
+## [0.2.0] - 2022-02-11
 ### Added
 - Overloads for benchmark parameters and functions to allow benchmark of custom classes.
 
 
-## [0.1.2] - 2022-02-10]
+## [0.1.2] - 2022-02-10
 ### Fixed
 - GPU compatibility.
 
 ### Removed
 - Carbon-tracker energy measurement. Library is still too immature at this point.
 
 
-## [0.1.1] - 2022-02-10]
+## [0.1.1] - 2022-02-10
 ### Added
 - Initial version.

pytorch_benchmark/benchmark.py

@@ -150,14 +150,29 @@ def measure_repeated_inference_timing(
     ):
         start_on_cpu = time()
         device_sample = transfer_to_device_fn(sample, model_device)
-        start_on_device = time()
+
+        if model_device.type == "cuda":
+            start_event = torch.cuda.Event(enable_timing=True)
+            stop_event = torch.cuda.Event(enable_timing=True)
+            start_event.record()  # For GPU timing
+        start_on_device = time()  # For CPU timing
+
         device_result = model(device_sample)
-        stop_on_device = time()
+
+        if model_device.type == "cuda":
+            stop_event.record()
+            torch.cuda.synchronize()
+            elapsed_on_device = stop_event.elapsed_time(start_event)
+            stop_on_device = time()
+        else:
+            stop_on_device = time()
+            elapsed_on_device = stop_on_device - start_on_device
+
         transfer_to_device_fn(device_result, "cpu")
         stop_on_cpu = time()
 
         t_c2d.append(start_on_device - start_on_cpu)
-        t_inf.append(stop_on_device - start_on_device)
+        t_inf.append(elapsed_on_device)
         t_d2c.append(stop_on_cpu - stop_on_device)
         t_tot.append(stop_on_cpu - start_on_cpu)
 

setup.py

@@ -25,7 +25,7 @@ def from_file(file_name: str = "requirements.txt", comment_char: str = "#"):
 
 setup(
     name="pytorch-benchmark",
-    version="0.2.0",
+    version="0.2.1",
     description="Easily benchmark PyTorch model FLOPs, latency, throughput, max allocated memory and energy consumption in one go.",
     long_description=long_description(),
     long_description_content_type="text/markdown",