fix: bn 线程共享错误
This commit is contained in:
+15
-2
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user