Skip to content

GitLab

Unverified Commit 08a9558f authored by Xiaomeng Zhao's avatar Xiaomeng Zhao Committed by GitHub
Browse files

Merge pull request #1594 from myhloli/dev 

perf(magic_pdf): adjust batch ratio calculation for GPU memory
parents 636d78a3 037736fb
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
magic_pdf/model/doc_analyze_by_custom_model.py
View file @ 08a9558f
...@@ -177,7 +177,7 @@ def doc_analyze( ...@@ -177,7 +177,7 @@ def doc_analyze(
if torch.cuda.is_available() and device != 'cpu' or npu_support: if torch.cuda.is_available() and device != 'cpu' or npu_support:
gpu_memory = int(os.getenv("VIRTUAL_VRAM_SIZE", round(get_vram(device)))) 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 is not None and gpu_memory >= 8:
batch_ratio = int(gpu_memory-5) batch_ratio = int(gpu_memory-6)
if batch_ratio >= 1: if batch_ratio >= 1:
logger.info(f'gpu_memory: {gpu_memory} GB, batch_ratio: {batch_ratio}') logger.info(f'gpu_memory: {gpu_memory} GB, batch_ratio: {batch_ratio}')
batch_model = BatchAnalyze(model=custom_model, batch_ratio=batch_ratio) batch_model = BatchAnalyze(model=custom_model, batch_ratio=batch_ratio)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment