"vscode:/vscode.git/clone" did not exist on "4874c96af18f4b3bf726f32c30a89c6fcd648fe9"
Commit e74a2960 authored by myhloli's avatar myhloli
Browse files

refactor(magic_pdf): adjust VRAM allocation and MFR batch size- Update VRAM...

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
parent 052a4d72
......@@ -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:
......
......@@ -89,7 +89,7 @@ class UnimernetModel(object):
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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment