Function test_gpu_optimizations

test_gitbase_e2e.py:98–176

Test GPU optimizations if available.

()

def test_gpu_optimizations():
    """Test GPU optimizations if available."""
    if not torch.cuda.is_available():
        print("⚠️  CUDA not available, skipping GPU tests")
        return

    print("=" * 60)
    print("TESTING GPU OPTIMIZATIONS")
    print("=" * 60)

    test_images = create_test_images()
    sample_image = list(test_images.values())[0]

    # Test different optimization combinations
    configs = [
        {"name": "Baseline FP32", "dtype": torch.float32, "use_compile": False, "use_channels_last": False},
        {"name": "FP16 Optimization", "dtype": torch.float16, "use_compile": False, "use_channels_last": False},
        {"name": "BF16 Optimization", "dtype": torch.bfloat16, "use_compile": False, "use_channels_last": False},
        {"name": "Channels Last", "dtype": torch.float16, "use_compile": False, "use_channels_last": True},
        {"name": "Torch Compile + FP16", "dtype": torch.float16, "use_compile": True, "use_channels_last": False},
        {"name": "All Optimizations", "dtype": torch.float16, "use_compile": True, "use_channels_last": True},
    ]

    results = {}

    for config in configs:
        try:
            print(f"\n--- Testing {config['name']} ---")

            captioner = GitBaseCaptioner(
                dtype=config["dtype"],
                device="cuda",
                use_compile=config["use_compile"],
                use_channels_last=config["use_channels_last"],
            )

            # Warmup run
            _ = captioner.caption_image_fast(sample_image)

            # Benchmark runs
            times = []
            captions = []

            for i in range(5):
                start_time = time.time()
                caption = captioner.caption_image_fast(sample_image)
                elapsed = time.time() - start_time
                times.append(elapsed)
                captions.append(caption)

            avg_time = sum(times) / len(times)
            results[config["name"]] = avg_time

            print(f"Average time: {avg_time:.3f}s")
            print(f"Sample caption: {captions[0]}")
            print(f"Caption consistency: {len(set(captions))} unique captions out of 5")

        except Exception as e: