Commit 4c9bf8ab authored by myhloli's avatar myhloli

refactor(memory management): remove unused clean_memory function

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.
parent 42a7d792
# 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()
\ No newline at end of file
...@@ -3,6 +3,7 @@ import os ...@@ -3,6 +3,7 @@ import os
import time import time
from magic_pdf.libs.Constants import * from magic_pdf.libs.Constants import *
from magic_pdf.libs.clean_memory import clean_memory
from magic_pdf.model.model_list import AtomicModel from magic_pdf.model.model_list import AtomicModel
os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新 os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新
...@@ -330,6 +331,8 @@ class CustomPEKModel: ...@@ -330,6 +331,8 @@ class CustomPEKModel:
elif int(res['category_id']) in [5]: elif int(res['category_id']) in [5]:
table_res_list.append(res) table_res_list.append(res)
clean_memory()
# ocr识别 # ocr识别
if self.apply_ocr: if self.apply_ocr:
ocr_start = time.time() ocr_start = time.time()
......
...@@ -7,6 +7,7 @@ from typing import List ...@@ -7,6 +7,7 @@ from typing import List
import torch 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.commons import fitz, get_delta_time
from magic_pdf.libs.convert_utils import dict_to_list from magic_pdf.libs.convert_utils import dict_to_list
from magic_pdf.libs.drop_reason import DropReason from magic_pdf.libs.drop_reason import DropReason
...@@ -304,14 +305,6 @@ def parse_page_core(pdf_docs, magic_model, page_id, pdf_bytes_md5, imageWriter, ...@@ -304,14 +305,6 @@ def parse_page_core(pdf_docs, magic_model, page_id, pdf_bytes_md5, imageWriter,
return page_info 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, def pdf_parse_union(pdf_bytes,
model_list, model_list,
imageWriter, imageWriter,
......
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