Forráskód Böngészése

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

- Simplify batch ratio logic for GPU memory >= 16GB
- Remove unnecessary conditions for 20GB and 40GB memory
myhloli 8 hónapja
szülő
commit
0d3304d7c9
1 módosított fájl, 1 hozzáadás és 5 törlés
  1. 1 5
      magic_pdf/model/doc_analyze_by_custom_model.py

+ 1 - 5
magic_pdf/model/doc_analyze_by_custom_model.py

@@ -170,11 +170,7 @@ def doc_analyze(
         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:
 
 
-            if gpu_memory >= 40:
-                batch_ratio = 32
-            elif gpu_memory >=20:
-                batch_ratio = 16
-            elif gpu_memory >= 16:
+            if gpu_memory >= 16:
                 batch_ratio = 8
                 batch_ratio = 8
             elif gpu_memory >= 10:
             elif gpu_memory >= 10:
                 batch_ratio = 4
                 batch_ratio = 4