YoloV5在Docker容器中加载模型时挂起问题求助

问题现象 #

本地Windows 10环境下，测试脚本可正常执行，输出finished后等待Ctrl+C终止；但在Docker容器中运行时，若脚本末尾包含无限循环，会卡在模型加载阶段的Adding AutoShape...步骤，无法继续执行；去掉末尾的无限循环则能正常完成运行。该问题导致基于watchdog.observers实现的文件监听推理系统无法正常工作。

测试代码 #

from time import sleep

import torch
import os
import pathlib


def start_watching():
    print("test no watchdog, yes loop")    

    try:        
        if os.name == 'nt':
            pathlib.PosixPath = pathlib.WindowsPath
        # 测试用硬编码路径
        model = torch.hub.load('./yolov5-master', 'custom', source='local', path='./best.pt', force_reload=True) 
        print("got model")
        results = model('./240206-154354_wheelset_149_22241687_Image_02.jpg')        
        labels, cord = results.xyxyn[0][:, -1], results.xyxyn[0][:, :-1]
        print(labels,cord)
    except Exception as ex:
        print(ex)
    
    print("sleepy time") 
    sleep(100)
    print("Finished")
    
    waitForFiles = True
    try:
        while waitForFiles:
            sleep(10)
    except KeyboardInterrupt:
        waitForFiles = False
        print("going down")


start_watching()

Docker运行输出 #

test no watchdog, yes loop
YOLOv5 🚀 2024-3-5 Python-3.11.6 torch-2.2.1+cpu CPU

Fusing layers...
Model summary: 157 layers, 7012822 parameters, 0 gradients, 15.8 GFLOPs
Adding AutoShape...

Dockerfile #

FROM yolobase
# 复制代码到容器
COPY . .
RUN echo "min not working"

ENTRYPOINT python3 test.py

环境说明 #

• 基础镜像yolobase基于YOLOv5官方Dockerfile构建，使用官方镜像ultralytics/yolov5:latest-cpu也存在相同问题
• 依赖配置如下：

# YOLOv5 requirements

# Usage: pip install -r requirements.txt

# Base ------------------------------------------------------------------------
gitpython>=3.1.30
matplotlib>=3.3
numpy>=1.23.5
opencv-python>=4.1.1
Pillow>=9.4.0
psutil  # system resources
PyYAML>=5.3.1
requests>=2.23.0
scipy>=1.4.1
thop>=0.1.1  # FLOPs computation
torch>=1.8.0  # see https://pytorch.org/get-started/locally (recommended)
torchvision>=0.9.0
tqdm>=4.64.0
ultralytics>=8.0.232
# protobuf<=3.20.1  # https://github.com/ultralytics/yolov5/issues/8012

# Logging ---------------------------------------------------------------------
# tensorboard>=2.4.1
# clearml>=1.2.0
# comet

# Plotting --------------------------------------------------------------------
pandas>=1.1.4
seaborn>=0.11.0

# Export ----------------------------------------------------------------------
# coremltools>=6.0  # CoreML export
# onnx>=1.10.0  # ONNX export
# onnx-simplifier>=0.4.1  # ONNX simplifier
# nvidia-pyindex  # TensorRT export
# nvidia-tensorrt  # TensorRT export
# scikit-learn<=1.1.2  # CoreML quantization
# tensorflow>=2.4.0,<=2.13.1  # TF exports (-cpu, -aarch64, -macos)
# tensorflowjs>=3.9.0  # TF.js export
# openvino-dev>=2023.0  # OpenVINO export

# Deploy ----------------------------------------------------------------------
setuptools>=65.5.1 # Snyk vulnerability fix
# tritonclient[all]~=2.24.0

# Extras ----------------------------------------------------------------------
# ipython  # interactive notebook
# mss  # screenshots
# albumentations>=1.0.3
# pycocotools>=2.0.6  # COCO mAP

解决方案尝试 #

1. 关闭Python输出缓冲：Docker容器中Python默认使用缓冲输出，可能导致日志未及时刷新，造成“假挂起”的错觉。修改Dockerfile的启动命令，添加-u参数关闭缓冲：

   ENTRYPOINT python3 -u test.py
   

2. 线程分离逻辑：YOLOv5的AutoShape加载过程可能涉及后台线程，与主线程的无限循环存在调度冲突。将模型加载和循环逻辑放到单独线程中执行：

   from time import sleep
   import threading
   import torch
   import os
   import pathlib
   
   
   def model_inference_loop():
       try:        
           if os.name == 'nt':
               pathlib.PosixPath = pathlib.WindowsPath
           # 测试用硬编码路径
           model = torch.hub.load('./yolov5-master', 'custom', source='local', path='./best.pt', force_reload=True) 
           print("got model")
           results = model('./240206-154354_wheelset_149_22241687_Image_02.jpg')        
           labels, cord = results.xyxyn[0][:, -1], results.xyxyn[0][:, :-1]
           print(labels,cord)
       except Exception as ex:
           print(ex)
       
       print("sleepy time") 
       sleep(100)
       print("Finished")
       
       waitForFiles = True
       try:
           while waitForFiles:
               sleep(10)
       except KeyboardInterrupt:
           waitForFiles = False
           print("going down")
   
   
   def start_watching():
       print("test no watchdog, yes loop")    
       t = threading.Thread(target=model_inference_loop)
       t.start()
       t.join()
   
   
   start_watching()
   

3. 增加容器资源分配：Docker默认资源限制可能导致模型加载缓慢或挂起，运行容器时指定更多CPU和内存：

   docker run --cpus 2 --memory 4g 你的镜像名称
   

内容的提问来源于stack exchange，提问作者ChrisUK