cosmetics

Author: NicolasHug

benchmarks/benchmark_transforms.py

@@ -28,19 +28,22 @@
 
 try:
     import cv2
+
     HAS_OPENCV = True
 except ImportError:
     HAS_OPENCV = False
 
 try:
     import albumentations as A
+
     HAS_ALBUMENTATIONS = True
 except ImportError:
     HAS_ALBUMENTATIONS = False
 
 try:
     import kornia as K
     import kornia.augmentation as KA
+
     HAS_KORNIA = True
 except ImportError:
     HAS_KORNIA = False
@@ -52,8 +55,6 @@
 NORM_STD = [0.229, 0.224, 0.225]
 
 
-
-
 def torchvision_pipeline(images: torch.Tensor, target_size: int) -> torch.Tensor:
     images = F.resize(
         images, size=(target_size, target_size), interpolation=F.InterpolationMode.BILINEAR, antialias=True
@@ -110,13 +111,15 @@ def kornia_pipeline(image: torch.Tensor, target_size: int) -> torch.Tensor:
 
 def run_benchmark(args) -> Dict[str, Any]:
     backend = args.backend.lower()
-    
+
     device = args.device.lower()
     # Check device compatibility
-    if device == 'cuda' and backend not in ['tv', 'tv-compiled']:
-        raise RuntimeError(f"CUDA device not supported for {backend} backend. Only 'tv' and 'tv-compiled' support CUDA.")
-    
-    if device == 'cuda' and not torch.cuda.is_available():
+    if device == "cuda" and backend not in ["tv", "tv-compiled"]:
+        raise RuntimeError(
+            f"CUDA device not supported for {backend} backend. Only 'tv' and 'tv-compiled' support CUDA."
+        )
+
+    if device == "cuda" and not torch.cuda.is_available():
         raise RuntimeError("CUDA not available. Install cuda-enabled torch and torchvision, or use 'cpu' device.")
 
     if backend == "opencv" and not HAS_OPENCV:
@@ -161,9 +164,9 @@ def generate_test_images():
         memory_format = torch.channels_last if args.contiguity == "CL" else torch.contiguous_format
         if memory_format == torch.channels_last:
             images = images.to(memory_format=torch.channels_last)
-        
+
         # Move to device for torchvision backends
-        if backend in ['tv', 'tv-compiled']:
+        if backend in ["tv", "tv-compiled"]:
             images = images.to(device)
 
         if args.batch_size == 1:
@@ -203,8 +206,6 @@ def generate_test_images():
     return {"backend": args.backend, "stats": stats}
 
 
-
-
 def main():
     parser = argparse.ArgumentParser(description="Benchmark torchvision transforms")
     parser.add_argument("--num-exp", type=int, default=100, help="Number of experiments we average over")
@@ -233,7 +234,7 @@ def main():
         "--backend", type=str.lower, choices=all_backends + ["all"], default="all", help="Backend to benchmark"
     )
     parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output")
-    parser.add_argument("--device", type=str, default="cpu", help="Device to use: cpu, cuda, or gpu (default: cpu)")
+    parser.add_argument("--device", type=str, default="cpu", help="Device to use: cpu or cuda (default: cpu)")
 
     args = parser.parse_args()
 

benchmarks/utils.py

@@ -10,18 +10,21 @@
 
 try:
     import cv2
+
     HAS_OPENCV = True
 except ImportError:
     HAS_OPENCV = False
 
 try:
     import albumentations as A
+
     HAS_ALBUMENTATIONS = True
 except ImportError:
     HAS_ALBUMENTATIONS = False
 
 try:
     import kornia as K
+
     HAS_KORNIA = True
 except ImportError:
     HAS_KORNIA = False
@@ -108,17 +111,9 @@ def print_comparison_table(results: List[Dict[str, Any]]) -> None:
 def print_benchmark_info(args):
     """Print benchmark configuration and library versions."""
     device = args.device.lower()
-    if device in ['gpu', 'cuda']:
-        device = 'cuda'
-    else:
-        device = 'cpu'
-        
-    memory_format = 'channels_last' if args.contiguity == 'CL' else 'channels_first'
-    
-    print("=" * 80)
-    print("BENCHMARK CONFIGURATION")
-    print("=" * 80)
-    
+
+    memory_format = "channels_last" if args.contiguity == "CL" else "channels_first"
+
     # Collect configuration info
     config = [
         ["Device", device],
@@ -128,41 +123,18 @@ def print_benchmark_info(args):
         ["Experiments", f"{args.num_exp} (+ {args.warmup} warmup)"],
         ["Input → output size", f"{args.min_size}-{args.max_size} → {args.target_size}×{args.target_size}"],
     ]
-    
+
     print(tabulate(config, headers=["Parameter", "Value"], tablefmt="simple"))
     print()
-    
-    print("=" * 80)
-    print("LIBRARY VERSIONS")
-    print("=" * 80)
-    
+
     # Collect library versions
     versions = [
         ["PyTorch", torch.__version__],
         ["TorchVision", torchvision.__version__],
+        ["OpenCV", cv2.__version__ if HAS_OPENCV else "Not available"],
+        ["PIL/Pillow", getattr(Image, '__version__', "Version unavailable")],
+        ["Albumentations", A.__version__ if HAS_ALBUMENTATIONS else "Not available"],
+        ["Kornia", K.__version__ if HAS_KORNIA else "Not available"],
     ]
-    
-    if HAS_OPENCV:
-        versions.append(["OpenCV", cv2.__version__])
-    else:
-        versions.append(["OpenCV", "Not available"])
-        
-    # PIL version
-    try:
-        versions.append(["PIL/Pillow", Image.__version__])
-    except AttributeError:
-        versions.append(["PIL/Pillow", "Version unavailable"])
-    
-    if HAS_ALBUMENTATIONS:
-        versions.append(["Albumentations", A.__version__])
-    else:
-        versions.append(["Albumentations", "Not available"])
-        
-    if HAS_KORNIA:
-        versions.append(["Kornia", K.__version__])
-    else:
-        versions.append(["Kornia", "Not available"])
-    
+
     print(tabulate(versions, headers=["Library", "Version"], tablefmt="simple"))
-    print("=" * 80)
-    print()