Skip to content

Commit

Permalink
refactor(memory management): remove unused clean_memory function
Browse files Browse the repository at this point in the history
The clean_memory function has been removed from pdf_parse_union_core_v2.py due to it not being used.
This change streamlines the code and prevents potential confusion regarding its purpose.
  • Loading branch information
myhloli committed Sep 29, 2024
1 parent 42a7d79 commit 4c9bf8a
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 8 deletions.
10 changes: 10 additions & 0 deletions magic_pdf/libs/clean_memory.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
# Copyright (c) Opendatalab. All rights reserved.
import torch
import gc


def clean_memory():
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
gc.collect()
3 changes: 3 additions & 0 deletions magic_pdf/model/pdf_extract_kit.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import time

from magic_pdf.libs.Constants import *
from magic_pdf.libs.clean_memory import clean_memory
from magic_pdf.model.model_list import AtomicModel

os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新
Expand Down Expand Up @@ -330,6 +331,8 @@ def __call__(self, image):
elif int(res['category_id']) in [5]:
table_res_list.append(res)

clean_memory()

# ocr识别
if self.apply_ocr:
ocr_start = time.time()
Expand Down
9 changes: 1 addition & 8 deletions magic_pdf/pdf_parse_union_core_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import torch

from magic_pdf.libs.clean_memory import clean_memory
from magic_pdf.libs.commons import fitz, get_delta_time
from magic_pdf.libs.convert_utils import dict_to_list
from magic_pdf.libs.drop_reason import DropReason
Expand Down Expand Up @@ -304,14 +305,6 @@ def parse_page_core(pdf_docs, magic_model, page_id, pdf_bytes_md5, imageWriter,
return page_info


def clean_memory():
import gc
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
gc.collect()


def pdf_parse_union(pdf_bytes,
model_list,
imageWriter,
Expand Down

0 comments on commit 4c9bf8a

Please sign in to comment.