fix: bn 线程共享错误

This commit is contained in:
zzw
2026-09-10 10:32:43 +08:00
parent 6a96d59226
commit 6994fa8458
+15 -2
View File
@@ -16,11 +16,15 @@ import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
import logging
import threading
from typing import Dict, List, Optional, Tuple
import requests
from config.settings import API_HOST, API_CALLBACK
# YOLO 模型非线程安全(fuse/BN 会竞态报 AttributeError: bn),推理需串行
_yolo_infer_lock = threading.Lock()
# ==================== 日志配置 ====================
def setup_logging(verbose: bool = False):
level = logging.DEBUG if verbose else logging.INFO
@@ -59,8 +63,9 @@ def cv2_imwrite_chinese(file_path: str, img: np.ndarray) -> bool:
# ==================== YOLO 痤疮检测 ====================
def detect_acne(img: np.ndarray, image_path: str, model: YOLO, output_path: str) -> Tuple[bool, Dict, str]:
try:
# 直接全图检测
results = model(img, conf=0.2, verbose=False)
# 共享模型下串行推理,避免多线程 fuse/BN 竞态
with _yolo_infer_lock:
results = model(img, conf=0.2, verbose=False)
annotated_img = img.copy()
class_counts = {}
@@ -632,6 +637,14 @@ def main():
logger.info("加载 YOLO 痤疮模型...")
try:
model = YOLO(model_path)
# 主线程预 fuse,避免首批多线程推理时争抢 BN
try:
if hasattr(model, "fuse"):
model.fuse()
elif hasattr(model, "model") and hasattr(model.model, "fuse"):
model.model.fuse()
except Exception as fuse_err:
logger.warning(f"YOLO 预 fuse 跳过:{fuse_err}")
logger.info("痤疮模型加载成功")
except Exception as e:
logger.error(f"痤疮模型加载失败:{e}")