Просмотр исходного кода

fix(magic_pdf): limit batch ratio for GPU memory

- 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
myhloli 10 месяцев назад
Родитель
Сommit
db8be9745d
1 измененных файлов с 5 добавлено и 1 удалено
  1. 5 1
      magic_pdf/model/doc_analyze_by_custom_model.py

+ 5 - 1
magic_pdf/model/doc_analyze_by_custom_model.py

@@ -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)