From 037736fbafb716f8fb072f3ae23ab6b63eebde9d Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 21 Jan 2025 19:47:49 +0800 Subject: [PATCH] perf(magic_pdf): adjust batch ratio calculation for GPU memory - Reduce batch_ratio by 1 for better performance and stability - This change ensures more consistent memory usage when processing documents --- magic_pdf/model/doc_analyze_by_custom_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index 328c4d68..c81bcec0 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -177,7 +177,7 @@ def doc_analyze( if torch.cuda.is_available() and device != 'cpu' or npu_support: gpu_memory = int(os.getenv("VIRTUAL_VRAM_SIZE", round(get_vram(device)))) if gpu_memory is not None and gpu_memory >= 8: - batch_ratio = int(gpu_memory-5) + batch_ratio = int(gpu_memory-6) if batch_ratio >= 1: logger.info(f'gpu_memory: {gpu_memory} GB, batch_ratio: {batch_ratio}') batch_model = BatchAnalyze(model=custom_model, batch_ratio=batch_ratio)