ssds.pytorch

这是一个基于 PyTorch 和 Python 3 实现的单发多框检测器及其变体的仓库。该仓库安装简单，并提供了丰富的可视化方法。我们希望这个仓库能够帮助大家更好地理解 SSD 类模型，并轻松地训练和部署这些模型。

目前，该仓库包含以下功能：

• 多种 SSD 变体：SSD、FPN、BiFPN、YOLO 等。
• 多种基础网络：ResNet、RegNet、MobileNet 等。
• 可视化SSD 类模型的特征，帮助用户理解模型的设计与性能。
• 快速训练与推理：利用 NVIDIA Apex 和 DALI 加速训练，并支持将模型转换为 ONNX 或 TensorRT 格式以便部署。

本仓库依赖于 ODTK、Detectron 和 TensorFlow 对象检测 API 的工作成果。感谢他们的贡献。

注意 当前版本的预训练模型尚未完成，请查看 上一版本 以获取更丰富的预训练模型。

目录

• 安装
• 使用
• 性能与模型库
• 可视化
• 文档

安装

需求

• Python >= 3.7
• CUDA >= 10.0
• PyTorch >= 1.4

基本安装：

conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
git clone https://github.com/ShuangXieIrene/ssds.pytorch.git
cd ssds.pytorch
python setup.py clean -a install

并行训练所需的额外 Python 库

目前，NVIDIA DALI 和 Apex 尚未包含在 requirements.txt 中，需要手动安装。

• DALI

pip install --extra-index-url https://developer.download.nvidia.com/compute/redist/cuda/10.0 nvidia-dali

• apex

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

Docker

git clone https://github.com/ShuangXieIrene/ssds.pytorch.git
docker build -t ssds:local ./ssds.pytorch/
docker run --gpus all -it --rm -v /data:/data ssds:local

使用

0. 通过可视化检查配置文件

在 配置文件 中定义网络，并根据可视化后的锚框调整配置。

python -m ssds.utils.visualize -cfg experiments/cfgs/tests/test.yml

1. 训练

# 基础训练
python -m ssds.utils.train -cfg experiments/cfgs/tests/test.yml
# 并行训练
python -m torch.distributed.launch --nproc_per_node={num_gpus} -m ssds.utils.train_ddp -cfg experiments/cfgs/tests/test.yml

2. 评估

python -m ssds.utils.train -cfg experiments/cfgs/tests/test.yml -e

3. 导出为 ONNX 或 TRT 模型

python -m ssds.utils.export -cfg experiments/cfgs/tests/test.yml -c best_mAP.pth -h

性能

可视化

ssds.pytorch 快速上手指南

ssds.pytorch 是一个基于 PyTorch 实现的单发多框检测器（SSD）及其变种的开源仓库。它支持多种 SSD 变体（如 FPN, BiFPN, YOLO 等）和骨干网络（ResNet, RegNet, MobileNet 等），并提供丰富的可视化功能以辅助理解模型设计。该工具利用 NVIDIA Apex 和 DALI 加速训练，并支持导出为 ONNX 或 TensorRT 格式以便部署。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐)
• Python: >= 3.7
• CUDA: >= 10.0
• PyTorch: >= 1.4
• GPU: 支持 CUDA 的 NVIDIA 显卡

安装步骤

1. 基础安装

首先创建 Conda 环境并安装 PyTorch 依赖，然后克隆代码库并进行安装。

# 安装 PyTorch 及相关依赖 (根据实际 CUDA 版本调整 cudatoolkit)
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch

# 克隆项目
git clone https://github.com/ShuangXieIrene/ssds.pytorch.git
cd ssds.pytorch

# 清理并安装项目
python setup.py clean -a install

2. 可选：并行训练加速组件

若需使用多卡并行训练以获得更快的速度，建议手动安装 NVIDIA DALI 和 Apex。

安装 DALI:

pip install --extra-index-url https://developer.download.nvidia.com/compute/redist/cuda/10.0 nvidia-dali

安装 Apex:

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

3. Docker 方式（替代方案）

如果您偏好使用 Docker，可以直接构建镜像运行：

git clone https://github.com/ShuangXieIrene/ssds.pytorch.git
docker build -t ssds:local ./ssds.pytorch/
docker run --gpus all -it --rm -v /data:/data ssds:local

注意：当前版本的预训练模型尚在完善中。如需使用丰富的预训练模型，请参考 v0.3.1 版本。

基本使用

以下是最核心的操作流程，涵盖配置检查、训练、评估及模型导出。所有命令均基于配置文件 experiments/cfgs/tests/test.yml，使用前请根据实际需求修改该文件。

1. 可视化配置与锚框

在正式训练前，建议先可视化网络结构和锚框（Anchor Boxes），以确认配置是否合理。

python -m ssds.utils.visualize -cfg experiments/cfgs/tests/test.yml

2. 模型训练

单卡训练：

python -m ssds.utils.train -cfg experiments/cfgs/tests/test.yml

多卡并行训练： 将 {num_gpus} 替换为您可用的 GPU 数量。

python -m torch.distributed.launch --nproc_per_node={num_gpus} -m ssds.utils.train_ddp -cfg experiments/cfgs/tests/test.yml

3. 模型评估

使用 -e 参数对模型进行评估。

python -m ssds.utils.train -cfg experiments/cfgs/tests/test.yml -e

4. 模型导出 (ONNX / TensorRT)

将训练好的模型（例如 best_mAP.pth）导出为部署格式。

python -m ssds.utils.export -cfg experiments/cfgs/tests/test.yml -c best_mAP.pth -h