From 49d140c55b3a0e0edcf2c2d0e852de1923b69226 Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 21 Jan 2025 18:32:18 +0800 Subject: [PATCH 1/3] perf(model): adjust batch size for layout and formula detection - Reduce YOLO_LAYOUT_BASE_BATCH_SIZE from 4 to 1 - Simplify batch ratio calculation for formula detection - Remove unused conditional logic in batch ratio determination --- magic_pdf/model/batch_analyze.py | 8 +++++--- magic_pdf/model/doc_analyze_by_custom_model.py | 8 ++------ 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/magic_pdf/model/batch_analyze.py b/magic_pdf/model/batch_analyze.py index 496bb43e..c1e719dd 100644 --- a/magic_pdf/model/batch_analyze.py +++ b/magic_pdf/model/batch_analyze.py @@ -19,7 +19,7 @@ get_adjusted_mfdetrec_res, get_ocr_result_list) # from magic_pdf.operators.models import InferenceResult -YOLO_LAYOUT_BASE_BATCH_SIZE = 4 +YOLO_LAYOUT_BASE_BATCH_SIZE = 1 MFD_BASE_BATCH_SIZE = 1 MFR_BASE_BATCH_SIZE = 16 @@ -56,7 +56,8 @@ def __call__(self, images: list) -> list: layout_images.append(pil_img) images_layout_res += self.model.layout_model.batch_predict( - layout_images, self.batch_ratio * YOLO_LAYOUT_BASE_BATCH_SIZE + # layout_images, self.batch_ratio * YOLO_LAYOUT_BASE_BATCH_SIZE + layout_images, YOLO_LAYOUT_BASE_BATCH_SIZE ) for image_index, useful_list in modified_images: @@ -78,7 +79,8 @@ def __call__(self, images: list) -> list: # 公式检测 mfd_start_time = time.time() images_mfd_res = self.model.mfd_model.batch_predict( - images, self.batch_ratio * MFD_BASE_BATCH_SIZE + # images, self.batch_ratio * MFD_BASE_BATCH_SIZE + images, MFD_BASE_BATCH_SIZE ) logger.info( f'mfd time: {round(time.time() - mfd_start_time, 2)}, image num: {len(images)}' diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index 6b649418..6cc09696 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -176,12 +176,8 @@ 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 = 2 - if 8 < gpu_memory: - batch_ratio = 4 - + if gpu_memory is not None and gpu_memory >= 7.5: + batch_ratio = int((gpu_memory-5) // 1) 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) From 052a4d72ed37ceba16ec6ac9aa6785edab7b11b1 Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 21 Jan 2025 19:10:25 +0800 Subject: [PATCH 2/3] perf(magic_pdf): optimize batch ratio calculation for GPU - Update GPU memory check and batch ratio calculation logic - Add support for virtual VRAM size environment variable - Improve logging for GPU memory and batch ratio --- magic_pdf/model/doc_analyze_by_custom_model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index 6cc09696..d0bcb870 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -175,9 +175,9 @@ def doc_analyze( npu_support = True 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.5: - batch_ratio = int((gpu_memory-5) // 1) + gpu_memory = int(os.getenv("virtual_vram_size", round(get_vram(device)))) + if gpu_memory is not None and gpu_memory >= 8: + batch_ratio = int(gpu_memory-5) 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) From e74a2960fcb24cca06e4ba1a631c2368058c0e82 Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 21 Jan 2025 19:22:33 +0800 Subject: [PATCH 3/3] refactor(magic_pdf): adjust VRAM allocation and MFR batch size- Update VRAM allocation logic to use 'VIRTUAL_VRAM_SIZE' environment variable - Reduce MFR (Math Formula Recognition) batch size from 64 to 32 --- magic_pdf/model/doc_analyze_by_custom_model.py | 2 +- magic_pdf/model/sub_modules/mfr/unimernet/Unimernet.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index d0bcb870..328c4d68 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -175,7 +175,7 @@ def doc_analyze( npu_support = True if torch.cuda.is_available() and device != 'cpu' or npu_support: - 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: batch_ratio = int(gpu_memory-5) if batch_ratio >= 1: diff --git a/magic_pdf/model/sub_modules/mfr/unimernet/Unimernet.py b/magic_pdf/model/sub_modules/mfr/unimernet/Unimernet.py index 54e46c56..9eff1ccd 100644 --- a/magic_pdf/model/sub_modules/mfr/unimernet/Unimernet.py +++ b/magic_pdf/model/sub_modules/mfr/unimernet/Unimernet.py @@ -89,7 +89,7 @@ def predict(self, mfd_res, image): mf_image_list.append(bbox_img) dataset = MathDataset(mf_image_list, transform=self.mfr_transform) - dataloader = DataLoader(dataset, batch_size=64, num_workers=0) + dataloader = DataLoader(dataset, batch_size=32, num_workers=0) mfr_res = [] for mf_img in dataloader: mf_img = mf_img.to(self.device)