如何导出YOLO-World的ONNX模型

我们使用的是这个GitHub仓库：: 4826695a84ce370dad7fe4cb555a6c7162f3751d

1. 下载预训练的PTH文件

首先需要从 Hugging Face 下载预训练好的 PTH 文件，参考链接：yolo_world_v2_m_obj365v1_goldg_pretrain.pth

2. 准备环境

以下是我所使用的基础环境配置：

• 操作系统：Ubuntu 22.04
• CUDA版本：11.8

a. 创建虚拟环境

使用 virtualenv 创建虚拟环境：

virtualenv env
source env/bin/activate

b. 安装对应的 PyTorch 版本

请务必使用 torch==1.11.0 并且使用 CUDA 版本。以下是我的安装命令：

# CUDA 11.3
pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113

c. 安装 MMCV

安装 mmcv：

pip install openmim
mim install mmcv==2.0.0

d. 安装合适的 numpy 版本

我发现默认安装的 numpy 版本过高（2.*），卸载后安装 1.* 版本，例如 1.24.4：

pip uninstall numpy
pip install numpy==1.24.4

e. 安装其他依赖

安装 requirements 文件中列出的其他依赖包。其中，basic_requirements.txt 中的 #opencv-python-headless==4.2.0.34 被我禁用，因为这个版本在我的环境中无法安装成功。

f. 拷贝 mmyolo 源码

需要将 mmyolo 的源码拷贝到 yolo-world 目录下, 请注意使用版本v0.6.0：

git clone https://github.com/open-mmlab/mmyolo.git third_party

g. 拷贝配置文件

将此文件拷贝到对应的 data/coco/lvis/ 目录中：lvis_v1_minival_inserted_image_name.json

h. 配置 Hugging Face 代理

此模型导出时需要依赖运行 Hugging Face 中的 Clip 模型，而国内连接 Hugging Face 不太稳定，建议使用代理。可以通过以下命令设置代理：

export http_proxy=http://192.168.0.106:7890
export https_proxy=http://192.168.0.106:7890
export all_proxy=socks5://192.168.0.106:7890

这里的 IP 地址和端口可以根据你局域网中已支持 VPN 的 Windows 机器来设置。

3. 导出模型

使用以下命令来导出 ONNX 模型：

PYTHONPATH=./ python deploy/export_onnx.py configs/pretrain/yolo_world_v2_m_vlpan_bn_2e-3_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.py weights/yolo_world_v2_m_obj365v1_goldg_pretrain-c6237d5b.pth --simplify --opset 12 --without-bbox-decoder --without-nms

注意：这里的 --without-bbox-decoder 一定要有，否则模型无法成功导出为 ONNX 格式。

4. 推理模型

推理模型时可以使用以下命令：

PYTHONPATH=./ python deploy/onnx_demo.py --onnx-nms work_dirs/yolo_world_v2_m_obj365v1_goldg_pretrain-c6237d5b.onnx 2.jpg  data/texts/coco_class_texts.json

但我在推理时没有得到结果，目前还在排查可能的原因。

PS: 我这边的pip requirements.txt环境

addict==2.4.0
aiofiles==23.2.1
albucore==0.0.20
albumentations==1.4.21
aliyun-python-sdk-core==2.16.0
aliyun-python-sdk-kms==2.16.5
altair==5.5.0
annotated-types==0.7.0
anyio==4.6.2.post1
attrs==24.2.0
certifi==2024.8.30
cffi==1.17.1
charset-normalizer==3.4.0
click==8.1.7
colorama==0.4.6
coloredlogs==15.0.1
contourpy==1.3.1
crcmod==1.7
cryptography==44.0.0
cycler==0.12.1
Cython==3.0.11
defusedxml==0.7.1
eval_type_backport==0.2.0
exceptiongroup==1.2.2
fastapi==0.115.5
ffmpy==0.4.0
filelock==3.14.0
flatbuffers==24.3.25
fonttools==4.55.1
fsspec==2024.10.0
gradio==4.16.0
gradio_client==0.8.1
h11==0.14.0
httpcore==1.0.7
httpx==0.28.0
huggingface-hub==0.26.3
humanfriendly==10.0
idna==3.10
importlib_resources==6.4.5
Jinja2==3.1.4
jmespath==0.10.0
jsonschema==4.23.0
jsonschema-specifications==2024.10.1
kiwisolver==1.4.7
lvis @ git+https://github.com/lvis-dataset/lvis-api.git@7d7f07def11da91f8b2710ce352c62a78fd5a7ad
Markdown==3.7
markdown-it-py==3.0.0
MarkupSafe==2.1.5
matplotlib==3.9.3
mdurl==0.1.2
mmcv==2.0.0
mmdet==3.0.0
mmengine==0.10.3
mmyolo==0.6.0
model-index==0.1.11
mpmath==1.3.0
narwhals==1.15.1
numpy==1.24.4
onnx==1.17.0
onnxruntime==1.20.1
onnxsim==0.4.36
opencv-python==4.9.0.80
opencv-python-headless==4.10.0.84
opendatalab==0.0.10
openmim==0.3.9
openxlab==0.1.2
ordered-set==4.1.0
orjson==3.10.12
oss2==2.17.0
packaging==24.2
pandas==2.2.3
pillow==10.4.0
platformdirs==4.3.6
prettytable==3.12.0
protobuf==5.29.0
pycocotools==2.0.8
pycparser==2.22
pycryptodome==3.21.0
pydantic==2.10.2
pydantic_core==2.27.1
pydub==0.25.1
Pygments==2.18.0
pyparsing==3.2.0
python-dateutil==2.9.0.post0
python-multipart==0.0.19
pytz==2023.4
PyYAML==6.0.2
referencing==0.35.1
regex==2024.11.6
requests==2.28.2
rich==13.4.2
rpds-py==0.22.0
ruff==0.8.1
safetensors==0.4.5
scipy==1.14.1
semantic-version==2.10.0
shapely==2.0.6
shellingham==1.5.4
simsimd==6.2.1
six==1.16.0
sniffio==1.3.1
starlette==0.41.3
stringzilla==3.11.0
supervision==0.25.0
sympy==1.13.3
tabulate==0.9.0
termcolor==2.5.0
terminaltables==3.1.10
timm==0.6.13
tokenizers==0.20.3
tomli==2.2.1
tomlkit==0.12.0
torch==1.11.0+cu113
torchaudio==0.11.0
torchvision==0.12.0+cu113
tqdm==4.65.2
transformers==4.46.3
typer==0.14.0
typing_extensions==4.12.2
tzdata==2024.2
urllib3==1.26.20
uvicorn==0.32.1
wcwidth==0.2.13
websockets==11.0.3
yapf==0.43.0