Function test_memory_usage

test_gitbase_e2e.py:323–361

Test memory usage and cleanup.

()

def test_memory_usage():
    """Test memory usage and cleanup."""
    if not torch.cuda.is_available():
        print("⚠️  CUDA not available, skipping memory tests")
        return

    print("=" * 60)
    print("TESTING MEMORY USAGE")
    print("=" * 60)

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

    # Clear cache
    torch.cuda.empty_cache()
    initial_memory = torch.cuda.memory_allocated()

    print(f"Initial GPU memory: {initial_memory / 1024**2:.1f} MB")

    # Create captioner and run inference
    captioner = GitBaseCaptioner(dtype=torch.float16, device="cuda", use_compile=False, use_channels_last=True)

    # Load model
    captioner.load()
    after_load_memory = torch.cuda.memory_allocated()

    print(f"Memory after model load: {after_load_memory / 1024**2:.1f} MB")
    print(f"Model memory usage: {(after_load_memory - initial_memory) / 1024**2:.1f} MB")

    # Run inference
    for i in range(10):
        captioner.caption_image_fast(sample_image)
        if i == 0:
            after_first_inference = torch.cuda.memory_allocated()
            print(f"Memory after first inference: {after_first_inference / 1024**2:.1f} MB")

    final_memory = torch.cuda.memory_allocated()
    print(f"Final memory: {final_memory / 1024**2:.1f} MB")
    print(f"Memory growth during inference: {(final_memory - after_first_inference) / 1024**2:.1f} MB")


def main():