Procházet zdrojové kódy

Added cuda_is_available (#4725)

### What problem does this PR solve?

Added cuda_is_available

### Type of change

- [x] Refactoring
tags/v0.16.0
Zhichang Yu před 8 měsíci
rodič
revize
3411d0a2ce
Žádný účet není propojen s e-mailovou adresou tvůrce revize
4 změnil soubory, kde provedl 24 přidání a 7 odebrání
  1. 10
    2
      deepdoc/vision/ocr.py
  2. 11
    2
      deepdoc/vision/recognizer.py
  3. 1
    1
      pyproject.toml
  4. 2
    2
      uv.lock

+ 10
- 2
deepdoc/vision/ocr.py Zobrazit soubor

import math import math
import numpy as np import numpy as np
import cv2 import cv2
import torch
import onnxruntime as ort import onnxruntime as ort


from .postprocess import build_post_process from .postprocess import build_post_process
raise ValueError("not find model file path {}".format( raise ValueError("not find model file path {}".format(
model_file_path)) model_file_path))


def cuda_is_available():
try:
import torch
if torch.cuda.is_available():
return True
except Exception:
return False
return False

options = ort.SessionOptions() options = ort.SessionOptions()
options.enable_cpu_mem_arena = False options.enable_cpu_mem_arena = False
options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
# https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580 # https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580
# Shrink GPU memory after execution # Shrink GPU memory after execution
run_options = ort.RunOptions() run_options = ort.RunOptions()
if torch.cuda.is_available():
if cuda_is_available():
cuda_provider_options = { cuda_provider_options = {
"device_id": 0, # Use specific GPU "device_id": 0, # Use specific GPU
"gpu_mem_limit": 512 * 1024 * 1024, # Limit gpu memory "gpu_mem_limit": 512 * 1024 * 1024, # Limit gpu memory

+ 11
- 2
deepdoc/vision/recognizer.py Zobrazit soubor

import cv2 import cv2
from copy import deepcopy from copy import deepcopy


import torch
import onnxruntime as ort import onnxruntime as ort
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download


if not os.path.exists(model_file_path): if not os.path.exists(model_file_path):
raise ValueError("not find model file path {}".format( raise ValueError("not find model file path {}".format(
model_file_path)) model_file_path))

def cuda_is_available():
try:
import torch
if torch.cuda.is_available():
return True
except Exception:
return False
return False

# https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580 # https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580
# Shrink GPU memory after execution # Shrink GPU memory after execution
self.run_options = ort.RunOptions() self.run_options = ort.RunOptions()


if torch.cuda.is_available():
if cuda_is_available():
options = ort.SessionOptions() options = ort.SessionOptions()
options.enable_cpu_mem_arena = False options.enable_cpu_mem_arena = False
cuda_provider_options = { cuda_provider_options = {

+ 1
- 1
pyproject.toml Zobrazit soubor

"tencentcloud-sdk-python==3.0.1215", "tencentcloud-sdk-python==3.0.1215",
"tika==2.6.0", "tika==2.6.0",
"tiktoken==0.7.0", "tiktoken==0.7.0",
"torch>=2.5.0,<3.0.0",
"umap_learn==0.5.6", "umap_learn==0.5.6",
"vertexai==1.64.0", "vertexai==1.64.0",
"volcengine==1.0.146", "volcengine==1.0.146",
"fastembed>=0.3.6,<0.4.0; sys_platform == 'darwin' or platform_machine != 'x86_64'", "fastembed>=0.3.6,<0.4.0; sys_platform == 'darwin' or platform_machine != 'x86_64'",
"fastembed-gpu>=0.3.6,<0.4.0; sys_platform != 'darwin' and platform_machine == 'x86_64'", "fastembed-gpu>=0.3.6,<0.4.0; sys_platform != 'darwin' and platform_machine == 'x86_64'",
"flagembedding==1.2.10", "flagembedding==1.2.10",
"torch>=2.5.0,<3.0.0",
"transformers>=4.35.0,<5.0.0" "transformers>=4.35.0,<5.0.0"
] ]

+ 2
- 2
uv.lock Zobrazit soubor

{ name = "tencentcloud-sdk-python" }, { name = "tencentcloud-sdk-python" },
{ name = "tika" }, { name = "tika" },
{ name = "tiktoken" }, { name = "tiktoken" },
{ name = "torch" },
{ name = "umap-learn" }, { name = "umap-learn" },
{ name = "valkey" }, { name = "valkey" },
{ name = "vertexai" }, { name = "vertexai" },
{ name = "fastembed", marker = "platform_machine != 'x86_64' or sys_platform == 'darwin'" }, { name = "fastembed", marker = "platform_machine != 'x86_64' or sys_platform == 'darwin'" },
{ name = "fastembed-gpu", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, { name = "fastembed-gpu", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" },
{ name = "flagembedding" }, { name = "flagembedding" },
{ name = "torch" },
{ name = "transformers" }, { name = "transformers" },
] ]


{ name = "tencentcloud-sdk-python", specifier = "==3.0.1215" }, { name = "tencentcloud-sdk-python", specifier = "==3.0.1215" },
{ name = "tika", specifier = "==2.6.0" }, { name = "tika", specifier = "==2.6.0" },
{ name = "tiktoken", specifier = "==0.7.0" }, { name = "tiktoken", specifier = "==0.7.0" },
{ name = "torch", specifier = ">=2.5.0,<3.0.0" },
{ name = "torch", marker = "extra == 'full'", specifier = ">=2.5.0,<3.0.0" },
{ name = "transformers", marker = "extra == 'full'", specifier = ">=4.35.0,<5.0.0" }, { name = "transformers", marker = "extra == 'full'", specifier = ">=4.35.0,<5.0.0" },
{ name = "umap-learn", specifier = "==0.5.6" }, { name = "umap-learn", specifier = "==0.5.6" },
{ name = "valkey", specifier = "==6.0.2" }, { name = "valkey", specifier = "==6.0.2" },

Načítá se…
Zrušit
Uložit