227 lines
7.8 KiB
Python
227 lines
7.8 KiB
Python
"""PaddleOCR-VL-1.6 GPU 单图推理与 Benchmark。"""
|
||
|
||
import argparse
|
||
import json
|
||
import platform
|
||
import statistics
|
||
import sys
|
||
import time
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
GPU_DIR = Path(__file__).resolve().parent
|
||
PROJECT_ROOT = GPU_DIR.parent
|
||
DEFAULT_IMAGE = PROJECT_ROOT / "images" / "手写01.png"
|
||
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "benchmarks" / "gpu"
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(description="PaddleOCR-VL-1.6 GPU Benchmark")
|
||
parser.add_argument("--image", type=Path, default=DEFAULT_IMAGE, help="待识别图片")
|
||
parser.add_argument("--device-id", type=int, default=0, help="CUDA GPU 编号")
|
||
parser.add_argument("--warmup", type=int, default=1, help="预热轮数")
|
||
parser.add_argument("--rounds", type=int, default=3, help="正式测试轮数")
|
||
parser.add_argument(
|
||
"--output-dir",
|
||
type=Path,
|
||
default=DEFAULT_OUTPUT_DIR,
|
||
help="Benchmark JSON 输出目录",
|
||
)
|
||
parser.add_argument("--no-result", action="store_true", help="不在控制台输出 OCR 文本")
|
||
return parser.parse_args()
|
||
|
||
|
||
def validate_args(args: argparse.Namespace) -> None:
|
||
args.image = args.image.expanduser().resolve()
|
||
args.output_dir = args.output_dir.expanduser().resolve()
|
||
if not args.image.is_file():
|
||
raise ValueError(f"图片不存在: {args.image}")
|
||
if args.device_id < 0:
|
||
raise ValueError("--device-id 不能小于 0")
|
||
if args.warmup < 0:
|
||
raise ValueError("--warmup 不能小于 0")
|
||
if args.rounds < 1:
|
||
raise ValueError("--rounds 必须大于等于 1")
|
||
|
||
|
||
def configure_cuda(device_id: int):
|
||
try:
|
||
import paddle
|
||
except ImportError as exc:
|
||
raise RuntimeError("未安装 GPU 子项目依赖，请先运行 gpu/setup_env.py。") from exc
|
||
|
||
if not paddle.is_compiled_with_cuda():
|
||
raise RuntimeError(
|
||
"当前 PaddlePaddle 未编译 CUDA 支持。请确认安装的是 paddlepaddle-gpu，"
|
||
"且正在使用 gpu/.venv。"
|
||
)
|
||
|
||
try:
|
||
device_count = paddle.device.cuda.device_count()
|
||
except Exception as exc:
|
||
raise RuntimeError(f"无法查询 CUDA 设备: {exc}") from exc
|
||
|
||
if device_count < 1:
|
||
raise RuntimeError("未检测到 NVIDIA CUDA GPU；本脚本不会自动回退到 CPU。")
|
||
if device_id >= device_count:
|
||
raise RuntimeError(f"GPU {device_id} 不存在，当前仅检测到 {device_count} 个 CUDA 设备。")
|
||
|
||
device = f"gpu:{device_id}"
|
||
try:
|
||
paddle.set_device(device)
|
||
paddle.device.cuda.synchronize(device_id)
|
||
except Exception as exc:
|
||
raise RuntimeError(f"无法启用 {device}: {exc}") from exc
|
||
|
||
try:
|
||
device_name = paddle.device.cuda.get_device_name(device_id)
|
||
except Exception:
|
||
device_name = "unknown"
|
||
|
||
return paddle, device, device_name
|
||
|
||
|
||
def synchronize(paddle: Any, device_id: int) -> None:
|
||
paddle.device.cuda.synchronize(device_id)
|
||
|
||
|
||
def read_gpu_memory(paddle: Any, device_id: int) -> dict[str, float | None]:
|
||
stats: dict[str, float | None] = {
|
||
"allocated_mb": None,
|
||
"reserved_mb": None,
|
||
"max_allocated_mb": None,
|
||
"max_reserved_mb": None,
|
||
}
|
||
functions = {
|
||
"allocated_mb": "memory_allocated",
|
||
"reserved_mb": "memory_reserved",
|
||
"max_allocated_mb": "max_memory_allocated",
|
||
"max_reserved_mb": "max_memory_reserved",
|
||
}
|
||
for key, function_name in functions.items():
|
||
function = getattr(paddle.device.cuda, function_name, None)
|
||
if function is None:
|
||
continue
|
||
try:
|
||
stats[key] = round(float(function(device_id)) / (1024**2), 2)
|
||
except Exception:
|
||
pass
|
||
return stats
|
||
|
||
|
||
def result_summary(result: list[Any]) -> dict[str, Any]:
|
||
first = result[0]
|
||
blocks = first["parsing_res_list"]
|
||
return {
|
||
"width": first["width"],
|
||
"height": first["height"],
|
||
"layout_boxes": len(first["layout_det_res"]["boxes"]),
|
||
"parsed_blocks": len(blocks),
|
||
"non_empty_blocks": sum(bool(block.content.strip()) for block in blocks),
|
||
}
|
||
|
||
|
||
def print_ocr_result(result: list[Any]) -> None:
|
||
print("\n[OCR Result]")
|
||
for item in result:
|
||
for block in item["parsing_res_list"]:
|
||
if block.content.strip():
|
||
print(f"[{block.label}] {block.bbox}")
|
||
print(block.content)
|
||
print()
|
||
|
||
|
||
def main() -> int:
|
||
args = parse_args()
|
||
try:
|
||
validate_args(args)
|
||
paddle, device, device_name = configure_cuda(args.device_id)
|
||
except (ValueError, RuntimeError) as exc:
|
||
print(f"[ERROR] {exc}", file=sys.stderr)
|
||
return 1
|
||
|
||
from paddleocr import PaddleOCRVL
|
||
|
||
print("=" * 70)
|
||
print(f"Device: {device} ({device_name})")
|
||
print(f"PaddlePaddle: {paddle.__version__}")
|
||
print(f"Input image: {args.image}")
|
||
print(f"Warmup/Rounds: {args.warmup}/{args.rounds}")
|
||
print("=" * 70)
|
||
|
||
synchronize(paddle, args.device_id)
|
||
init_started = time.perf_counter()
|
||
pipeline = PaddleOCRVL(pipeline_version="v1.6", device=device)
|
||
synchronize(paddle, args.device_id)
|
||
init_seconds = time.perf_counter() - init_started
|
||
print(f"Model init: {init_seconds:.3f}s")
|
||
|
||
result = None
|
||
for index in range(args.warmup):
|
||
print(f"Warmup {index + 1}/{args.warmup}...", flush=True)
|
||
result = pipeline.predict(str(args.image))
|
||
synchronize(paddle, args.device_id)
|
||
|
||
inference_times: list[float] = []
|
||
for index in range(args.rounds):
|
||
synchronize(paddle, args.device_id)
|
||
started = time.perf_counter()
|
||
result = pipeline.predict(str(args.image))
|
||
synchronize(paddle, args.device_id)
|
||
elapsed = time.perf_counter() - started
|
||
inference_times.append(elapsed)
|
||
print(f"Inference {index + 1}/{args.rounds}: {elapsed:.3f}s", flush=True)
|
||
|
||
if result is None:
|
||
print("[ERROR] 未产生推理结果。", file=sys.stderr)
|
||
return 2
|
||
|
||
summary = result_summary(result)
|
||
benchmark = {
|
||
"status": "completed",
|
||
"timestamp": datetime.now().astimezone().isoformat(),
|
||
"platform": platform.platform(),
|
||
"python_version": platform.python_version(),
|
||
"paddle_version": paddle.__version__,
|
||
"pipeline_version": "v1.6",
|
||
"device": device,
|
||
"device_name": device_name,
|
||
"image_path": str(args.image),
|
||
"image": summary,
|
||
"warmup_rounds": args.warmup,
|
||
"benchmark_rounds": args.rounds,
|
||
"model_init_seconds": round(init_seconds, 3),
|
||
"inference_seconds": {
|
||
"all": [round(value, 3) for value in inference_times],
|
||
"min": round(min(inference_times), 3),
|
||
"max": round(max(inference_times), 3),
|
||
"mean": round(statistics.fmean(inference_times), 3),
|
||
"median": round(statistics.median(inference_times), 3),
|
||
"stdev": round(statistics.pstdev(inference_times), 3),
|
||
},
|
||
"gpu_memory": read_gpu_memory(paddle, args.device_id),
|
||
}
|
||
|
||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
output_path = args.output_dir / f"gpu-benchmark-{timestamp}.json"
|
||
output_path.write_text(json.dumps(benchmark, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
print("\n[Benchmark]")
|
||
print(f"Image: {summary['width']} x {summary['height']}")
|
||
print(f"Layout boxes: {summary['layout_boxes']}")
|
||
print(f"Parsed blocks: {summary['parsed_blocks']}")
|
||
print(f"Average: {benchmark['inference_seconds']['mean']:.3f}s")
|
||
print(f"Min/Max: {benchmark['inference_seconds']['min']:.3f}s / {benchmark['inference_seconds']['max']:.3f}s")
|
||
print(f"Result JSON: {output_path}")
|
||
|
||
if not args.no_result:
|
||
print_ocr_result(result)
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|