Procházet zdrojové kódy

Merge pull request #1824 from myhloli/dev

perf(inference): adjust batch ratio for high GPU memory
Xiaomeng Zhao před 8 měsíci
rodič
revize
d327e1e996
1 změnil soubory, kde provedl 3 přidání a 1 odebrání
  1. 3 1
      magic_pdf/model/doc_analyze_by_custom_model.py

+ 3 - 1
magic_pdf/model/doc_analyze_by_custom_model.py

@@ -170,7 +170,9 @@ 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 >= 10:
+            if gpu_memory >= 16:
+                batch_ratio = 8
+            elif gpu_memory >= 10:
                 batch_ratio = 4
                 batch_ratio = 4
             else:
             else:
                 batch_ratio = 2
                 batch_ratio = 2