292 lines
11 KiB
Python
292 lines
11 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
|
||
if str(PROJECT_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(PROJECT_ROOT))
|
||
|
||
from ocr_logging import default_log_path, setup_run_logger
|
||
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 文本")
|
||
parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径")
|
||
parser.add_argument("--verbose", action="store_true", help="输出详细日志")
|
||
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:
|
||
program_started = time.perf_counter()
|
||
args = parse_args()
|
||
log_file = args.log_file or default_log_path(
|
||
PROJECT_ROOT,
|
||
"single",
|
||
args.image.stem,
|
||
device=f"gpu{args.device_id}",
|
||
)
|
||
logger = setup_run_logger("ocr.single.gpu", log_file, verbose=args.verbose)
|
||
logger.info(
|
||
"PROGRAM_STARTED image=%s device_id=%d warmup=%d rounds=%d output_dir=%s",
|
||
args.image,
|
||
args.device_id,
|
||
args.warmup,
|
||
args.rounds,
|
||
args.output_dir,
|
||
)
|
||
try:
|
||
validate_args(args)
|
||
cuda_started = time.perf_counter()
|
||
paddle, device, device_name = configure_cuda(args.device_id)
|
||
cuda_setup_seconds = time.perf_counter() - cuda_started
|
||
except (ValueError, RuntimeError) as exc:
|
||
logger.error("VALIDATION_OR_CUDA_FAILED type=%s error=%s", type(exc).__name__, exc, exc_info=args.verbose)
|
||
return 1
|
||
|
||
import_started = time.perf_counter()
|
||
from paddleocr import PaddleOCRVL
|
||
import_seconds = time.perf_counter() - import_started
|
||
logger.info(
|
||
"RUNTIME_READY cuda_setup_seconds=%.3f import_seconds=%.3f device=%s device_name=%s paddle_version=%s image_size_bytes=%d",
|
||
cuda_setup_seconds,
|
||
import_seconds,
|
||
device,
|
||
device_name,
|
||
paddle.__version__,
|
||
args.image.stat().st_size,
|
||
)
|
||
|
||
logger.info("MODEL_INITIALIZATION_STARTED pipeline_version=v1.6 device=%s", device)
|
||
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
|
||
logger.info("MODEL_INITIALIZED seconds=%.3f", init_seconds)
|
||
|
||
result = None
|
||
warmup_times: list[float] = []
|
||
for index in range(args.warmup):
|
||
started = time.perf_counter()
|
||
result = pipeline.predict(str(args.image))
|
||
synchronize(paddle, args.device_id)
|
||
elapsed = time.perf_counter() - started
|
||
warmup_times.append(elapsed)
|
||
logger.info("WARMUP_COMPLETED round=%d/%d seconds=%.3f", index + 1, args.warmup, elapsed)
|
||
|
||
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)
|
||
logger.info("INFERENCE_COMPLETED round=%d/%d seconds=%.3f", index + 1, args.rounds, elapsed)
|
||
|
||
if result is None:
|
||
logger.error("EMPTY_RESULT")
|
||
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,
|
||
"cuda_setup_seconds": round(cuda_setup_seconds, 3),
|
||
"runtime_import_seconds": round(import_seconds, 3),
|
||
"model_init_seconds": round(init_seconds, 3),
|
||
"warmup_seconds": [round(value, 3) for value in warmup_times],
|
||
"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),
|
||
"program_total_seconds": round(time.perf_counter() - program_started, 3),
|
||
"log_file": str(log_file.resolve()),
|
||
}
|
||
|
||
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")
|
||
|
||
logger.info(
|
||
"RESULT_SUMMARY width=%d height=%d layout_boxes=%d parsed_blocks=%d non_empty_blocks=%d gpu_memory=%s",
|
||
summary["width"],
|
||
summary["height"],
|
||
summary["layout_boxes"],
|
||
summary["parsed_blocks"],
|
||
summary["non_empty_blocks"],
|
||
benchmark["gpu_memory"],
|
||
)
|
||
logger.info(
|
||
"BENCHMARK_SUMMARY cuda_setup_seconds=%.3f import_seconds=%.3f model_init_seconds=%.3f warmup_total_seconds=%.3f inference_total_seconds=%.3f inference_min_seconds=%.3f inference_max_seconds=%.3f inference_mean_seconds=%.3f inference_median_seconds=%.3f inference_stdev_seconds=%.3f program_total_seconds=%.3f result_json=%s log=%s",
|
||
cuda_setup_seconds,
|
||
import_seconds,
|
||
init_seconds,
|
||
sum(warmup_times),
|
||
sum(inference_times),
|
||
min(inference_times),
|
||
max(inference_times),
|
||
statistics.fmean(inference_times),
|
||
statistics.median(inference_times),
|
||
statistics.pstdev(inference_times),
|
||
time.perf_counter() - program_started,
|
||
output_path,
|
||
log_file.resolve(),
|
||
)
|
||
|
||
if not args.no_result:
|
||
for index, block in enumerate(result[0]["parsing_res_list"], start=1):
|
||
if block.content.strip():
|
||
logger.info(
|
||
"OCR_BLOCK index=%d label=%s bbox=%s content=%s",
|
||
index,
|
||
block.label,
|
||
block.bbox,
|
||
block.content.replace("\r", "").replace("\n", "\\n"),
|
||
)
|
||
|
||
logger.info("PROGRAM_COMPLETED")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|