fix: bn 线程共享错误
This commit is contained in:
+14
-1
@@ -16,11 +16,15 @@ import argparse
|
|||||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
import logging
|
import logging
|
||||||
|
import threading
|
||||||
from typing import Dict, List, Optional, Tuple
|
from typing import Dict, List, Optional, Tuple
|
||||||
import requests
|
import requests
|
||||||
|
|
||||||
from config.settings import API_HOST, API_CALLBACK
|
from config.settings import API_HOST, API_CALLBACK
|
||||||
|
|
||||||
|
# YOLO 模型非线程安全(fuse/BN 会竞态报 AttributeError: bn),推理需串行
|
||||||
|
_yolo_infer_lock = threading.Lock()
|
||||||
|
|
||||||
# ==================== 日志配置 ====================
|
# ==================== 日志配置 ====================
|
||||||
def setup_logging(verbose: bool = False):
|
def setup_logging(verbose: bool = False):
|
||||||
level = logging.DEBUG if verbose else logging.INFO
|
level = logging.DEBUG if verbose else logging.INFO
|
||||||
@@ -59,7 +63,8 @@ def cv2_imwrite_chinese(file_path: str, img: np.ndarray) -> bool:
|
|||||||
# ==================== YOLO 痤疮检测 ====================
|
# ==================== YOLO 痤疮检测 ====================
|
||||||
def detect_acne(img: np.ndarray, image_path: str, model: YOLO, output_path: str) -> Tuple[bool, Dict, str]:
|
def detect_acne(img: np.ndarray, image_path: str, model: YOLO, output_path: str) -> Tuple[bool, Dict, str]:
|
||||||
try:
|
try:
|
||||||
# 直接全图检测
|
# 共享模型下串行推理,避免多线程 fuse/BN 竞态
|
||||||
|
with _yolo_infer_lock:
|
||||||
results = model(img, conf=0.2, verbose=False)
|
results = model(img, conf=0.2, verbose=False)
|
||||||
|
|
||||||
annotated_img = img.copy()
|
annotated_img = img.copy()
|
||||||
@@ -632,6 +637,14 @@ def main():
|
|||||||
logger.info("加载 YOLO 痤疮模型...")
|
logger.info("加载 YOLO 痤疮模型...")
|
||||||
try:
|
try:
|
||||||
model = YOLO(model_path)
|
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("痤疮模型加载成功")
|
logger.info("痤疮模型加载成功")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"痤疮模型加载失败:{e}")
|
logger.error(f"痤疮模型加载失败:{e}")
|
||||||
|
|||||||
Reference in New Issue
Block a user