Skip to content

Commit

Permalink
fix(magic_pdf): limit batch ratio for GPU memory
Browse files Browse the repository at this point in the history
- Commented out the original batch ratio calculation
- Set a fixed batch ratio of 2 for GPUs with less than 8 GB memory
- Increased batch ratio to 4 for GPUs with 8 GB or more memory
  • Loading branch information
myhloli committed Jan 17, 2025
1 parent 59502e5 commit db8be97
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 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,11 @@ def doc_analyze(
if torch.cuda.is_available() and device != 'cpu' or npu_support:
gpu_memory = get_vram(device)
if gpu_memory is not None and gpu_memory >= 7:
batch_ratio = int((gpu_memory-3) // 1.5)
# batch_ratio = int((gpu_memory-3) // 1.5)
batch_ratio = 2
if 8 < gpu_memory:
batch_ratio = 4

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 db8be97

Please sign in to comment.