Quellcode durchsuchen

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

- Remove separate condition for GPU memory >= 24GB
- Simplify logic to use a single threshold of 16GB
myhloli vor 8 Monaten
Ursprung
Commit
2f40fa7d3c
1 geänderte Dateien mit 1 neuen und 3 gelöschten Zeilen
  1. 1 3
      magic_pdf/model/doc_analyze_by_custom_model.py

+ 1 - 3
magic_pdf/model/doc_analyze_by_custom_model.py

@@ -257,9 +257,7 @@ 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 >= 24:
-                batch_ratio = 32
-            elif gpu_memory >= 16:
+            if gpu_memory >= 16:
                 batch_ratio = 16
             elif gpu_memory >= 12:
                 batch_ratio = 8