Skip to content

Commit

Permalink
perf(magic_pdf): adjust batch ratio calculation for GPU memory
Browse files Browse the repository at this point in the history
- Reduce batch_ratio by 1 for better performance and stability
- This change ensures more consistent memory usage when processing documents
  • Loading branch information
myhloli committed Jan 21, 2025
1 parent e74a296 commit 037736f
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion magic_pdf/model/doc_analyze_by_custom_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 037736f

Please sign in to comment.