Selaa lähdekoodia

perf(inference): adjust batch ratio thresholds for GPU memory sizes

- Increase batch ratio to 32 for GPU memory >= 24GB
- Set batch ratio to 16 for GPU memory >= 16GB
- Reduce batch ratio to 8 for GPU memory >= 12GB
- Lower batch ratio to 4 for GPU memory >= 8GB
- Set batch ratio to 2 for GPU memory >= 6GB
- Keep batch ratio at 1 for lower GPU memory sizes
myhloli 8 kuukautta sitten
vanhempi
commit
74e954dace
1 muutettua tiedostoa jossa 6 lisäystä ja 4 poistoa
  1. 6 4
      magic_pdf/model/doc_analyze_by_custom_model.py

+ 6 - 4
magic_pdf/model/doc_analyze_by_custom_model.py

@@ -257,13 +257,15 @@ def may_batch_image_analyze(
     if str(device).startswith('npu') or str(device).startswith('cuda'):
         gpu_memory = int(os.getenv('VIRTUAL_VRAM_SIZE', round(get_vram(device))))
         if gpu_memory is not None:
-            if gpu_memory >= 20:
+            if gpu_memory >= 24:
+                batch_ratio = 32
+            elif gpu_memory >= 16:
                 batch_ratio = 16
-            elif gpu_memory >= 15:
+            elif gpu_memory >= 12:
                 batch_ratio = 8
-            elif gpu_memory >= 10:
+            elif gpu_memory >= 8:
                 batch_ratio = 4
-            elif gpu_memory >= 7:
+            elif gpu_memory >= 6:
                 batch_ratio = 2
             else:
                 batch_ratio = 1