Merge pull request #1824 from myhloli/dev

perf(inference): adjust batch ratio for high GPU memory

Merge pull request #1824 from myhloli/dev
perf(inference): adjust batch ratio for high GPU memory
d327e1e9 · Xiaomeng Zhao · GitHub · f20ab37d · 0b05dff7 · d327e1e9
Unverified Commit d327e1e9 authored Mar 03, 2025 by Xiaomeng Zhao Committed by GitHub Mar 03, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

magic_pdf/model/doc_analyze_by_custom_model.py magic_pdf/model/doc_analyze_by_custom_model.py +3 -1

No files found.
--- a/magic_pdf/model/doc_analyze_by_custom_model.py
+++ b/magic_pdf/model/doc_analyze_by_custom_model.py
@@ -170,7 +170,9 @@ def doc_analyze(
        gpu_memory = int(os.getenv("VIRTUAL_VRAM_SIZE", round(get_vram(device))))
        if gpu_memory is not None and gpu_memory >= 8:

-            if gpu_memory >= 10:
+            if gpu_memory >= 16:
+                batch_ratio = 8
+            elif gpu_memory >= 10:
                batch_ratio = 4
            else:
                batch_ratio = 2