PyTorch-YOLOv4

YOLOv4 的一个极简 PyTorch 实现。

• YOLO v4 论文：https://arxiv.org/abs/2004.10934

• 源代码：https://github.com/AlexeyAB/darknet

• 更多详情：http://pjreddie.com/darknet/yolo/

• 推理

• 训练

  • 马赛克数据增强

├── README.md
├── dataset.py            数据集
├── demo.py               运行 PyTorch 的示例 --> 工具/darknet2pytorch
├── demo_darknet2onnx.py  转换为 ONNX 的工具 --> 工具/darknet2pytorch
├── demo_pytorch2onnx.py  转换为 ONNX 的工具
├── models.py             PyTorch 模型
├── train.py              训练 models.py
├── cfg.py                训练用的 cfg 文件
├── cfg                   cfg 文件 --> darknet2pytorch
├── data            
├── weight                --> darknet2pytorch
├── tool
│   ├── camera.py           一个摄像头示例
│   ├── coco_annotation.py       COCO 数据集生成器
│   ├── config.py
│   ├── darknet2pytorch.py
│   ├── region_loss.py
│   ├── utils.py
│   └── yolo_layer.py

0. 权重下载

0.1 Darknet

• 百度网盘(https://pan.baidu.com/s/1dAGEW8cm-dqK14TbhhVetA 提取码：dm5b)
• Google Drive(https://drive.google.com/open?id=1cewMfusmPjYWbrnuJRuKhPMwRe_b9PaT)

0.2 PyTorch

你可以使用 darknet2pytorch 自行转换，或者直接下载我转换好的模型。

• 百度网盘
  • yolov4.pth(https://pan.baidu.com/s/1ZroDvoGScDgtE1ja_QqJVw 提取码：xrq9)
  • yolov4.conv.137.pth(https://pan.baidu.com/s/1ovBie4YyVQQoUrC3AY0joA 提取码：kcel)
• Google Drive
  • yolov4.pth(https://drive.google.com/open?id=1wv_LiFeCRYwtpkqREPeI13-gPELBDwuJ)
  • yolov4.conv.137.pth(https://drive.google.com/open?id=1fcbR0bWzYfIEdLJPzOsn4R5mlvR6IQyA)

1. 训练

使用 YOLOv4 训练自己的数据

1. 下载权重

2. 转换数据

   对于 COCO 数据集，可以使用 tool/coco_annotation.py。

   # train.txt
   图片路径1 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
   图片路径2 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
   ...
   ...
   

3. 训练

   你可以在 cfg.py 中设置参数。

    python train.py -g [GPU_ID] -dir [数据集路径] ...
   

2. 推理

2.1 在 MS COCO 数据集上的性能（使用来自 https://github.com/AlexeyAB/darknet 的预训练 Darknet 权重）

ONNX 和 TensorRT 模型是从 PyTorch 转换而来（TianXiaomo）：PyTorch->ONNX->TensorRT。 更多关于转换的细节请参见后续章节。

• val2017 数据集（输入尺寸：416x416）

模型类型	AP	AP50	AP75	APS	APM	APL
DarkNet (YOLOv4 论文)	0.471	0.710	0.510	0.278	0.525	0.636
PyTorch (TianXiaomo)	0.466	0.704	0.505	0.267	0.524	0.629
TensorRT FP32 + BatchedNMSPlugin	0.472	0.708	0.511	0.273	0.530	0.637
TensorRT FP16 + BatchedNMSPlugin	0.472	0.708	0.511	0.273	0.530	0.636

• testdev2017 数据集（输入尺寸：416x416）

模型类型	AP	AP50	AP75	APS	APM	APL
DarkNet (YOLOv4 论文)	0.412	0.628	0.443	0.204	0.444	0.560
PyTorch (TianXiaomo)	0.404	0.615	0.436	0.196	0.438	0.552
TensorRT FP32 + BatchedNMSPlugin	0.412	0.625	0.445	0.200	0.446	0.564
TensorRT FP16 + BatchedNMSPlugin	0.412	0.625	0.445	0.200	0.446	0.563

2.2 推理时的图像输入尺寸

图像输入尺寸并不局限于 320 * 320、416 * 416、512 * 512 和 608 * 608。 你可以根据需要调整输入尺寸以获得不同的宽高比，例如：320 * 608。 更大的输入尺寸有助于检测更小的目标，但可能会降低推理速度并消耗更多 GPU 内存。

height = 320 + 96 * n, n in {0, 1, 2, 3, ...}
width  = 320 + 96 * m, m in {0, 1, 2, 3, ...}

2.3 不同的推理选项

• 加载预训练的 Darknet 模型和权重进行推理（图像尺寸已在 cfg 文件中配置）

  python demo.py -cfgfile <cfgFile> -weightfile <weightFile> -imgfile <imgFile>
  

• 加载 PyTorch 权重（pth 文件）进行推理

  python models.py <num_classes> <weightfile> <imgfile> <IN_IMAGE_H> <IN_IMAGE_W> <namefile(可选)>
  

• 加载转换后的 ONNX 文件进行推理（见第 3 和第 4 节）

• 加载转换后的 TensorRT 引擎文件进行推理（见第 5 节）

2.4 推理输出

推理输出有两个：

• 一个是边界框的位置信息，其形状为 [batch, num_boxes, 1, 4]，表示每个边界框的 x1, y1, x2, y2。
• 另一个是边界框的得分，形状为 [batch, num_boxes, num_classes]，表示每个边界框在所有类别上的得分。

目前仍需要进行一些后处理，包括 NMS。我们正在努力减少后处理的时间和复杂性。

3. Darknet2ONNX

• 此脚本用于将官方预训练的 Darknet 模型转换为 ONNX 格式

• 推荐的 PyTorch 版本：

  • 对于 TensorRT 7.0 及以上版本，建议使用 PyTorch 1.4.0
  • 对于 TensorRT 7.1.2 及以上版本，建议使用 PyTorch 1.5.0 或 1.6.0

• 安装 onnxruntime

  pip install onnxruntime
  

• 运行 Python 脚本生成 ONNX 模型并运行示例

  python demo_darknet2onnx.py <cfgFile> <namesFile> <weightFile> <imageFile> <batchSize>
  

3.1 动态或静态批次大小

• 正数批次大小会生成静态批次大小的 ONNX 模型，否则批次大小将是动态的
  • 动态批次大小只会生成一个 ONNX 模型
  • 静态批次大小会生成两个 ONNX 模型，其中一个用于运行示例（batch_size=1）

4. PyTorch转ONNX

• 您可以使用此脚本将训练好的PyTorch模型转换为ONNX格式

• 推荐的PyTorch版本：

  • 对于TensorRT 7.0及以上版本，建议使用PyTorch 1.4.0
  • 对于TensorRT 7.1.2及以上版本，建议使用PyTorch 1.5.0或1.6.0

• 安装onnxruntime

  pip install onnxruntime
  

• 运行Python脚本生成ONNX模型并执行演示

  python demo_pytorch2onnx.py <weight_file> <image_path> <batch_size> <n_classes> <IN_IMAGE_H> <IN_IMAGE_W>
  

  例如：

  python demo_pytorch2onnx.py yolov4.pth dog.jpg 8 80 416 416
  

4.1 动态或静态批大小

• 正数批大小将生成静态批大小的ONNX模型，否则批大小将是动态的
  • 动态批大小只会生成一个ONNX模型
  • 静态批大小会生成两个ONNX模型，其中一个用于运行演示（批大小=1）

5. ONNX转TensorRT

• 推荐的TensorRT版本：7.0、7.1

5.1 从静态批大小的ONNX转换

• 运行以下命令将YOLOv4 ONNX模型转换为TensorRT引擎

  trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
  

  • 注意：如果要在转换中使用int8模式，需要额外进行int8校准。

5.2 从动态批大小的ONNX转换

• 运行以下命令将YOLOv4 ONNX模型转换为TensorRT引擎

  trtexec --onnx=<onnx_file> \
  --minShapes=input:<shape_of_min_batch> --optShapes=input:<shape_of_opt_batch> --maxShapes=input:<shape_of_max_batch> \
  --workspace=<size_in_megabytes> --saveEngine=<engine_file> --fp16
  

• 例如：

  trtexec --onnx=yolov4_-1_3_320_512_dynamic.onnx \
  --minShapes=input:1x3x320x512 --optShapes=input:4x3x320x512 --maxShapes=input:8x3x320x512 \
  --workspace=2048 --saveEngine=yolov4_-1_3_320_512_dynamic.engine --fp16
  

5.3 运行演示

python demo_trt.py <tensorRT_engine_file> <input_image> <input_H> <input_W>

• 此处的演示仅在批大小为动态（1应在动态范围内）或批大小为1时有效，但您可以稍作修改以适应其他动态或静态批大小。

• 注意1：input_H和input_W应与原始ONNX文件中的输入尺寸一致。

• 注意2：TensorRT输出需要额外的NMS操作。此演示使用tool/utils.py中的Python NMS代码。

6. ONNX转TensorFlow

• 第一步：转换为ONNX

  tensorflow >=2.0

  1：感谢：github:https://github.com/onnx/onnx-tensorflow

  2：运行git clone https://github.com/onnx/onnx-tensorflow.git && cd onnx-tensorflow 然后运行pip install -e .

  注意：使用“pip install onnx-tf”时可能会出现错误，至少对我而言如此，因此建议使用源码安装方法。

7. ONNX转TensorRT及DeepStream推理

1. 编译DeepStream Nvinfer插件

      cd DeepStream
      make
   

2. 构建TRT引擎。

 对于单批次，

trtexec --onnx= --explicitBatch --saveEngine= --workspace= --fp16

对于多批次，

trtexec --onnx= --explicitBatch --shapes=input:Xx3xHxW --optShapes=input:Xx3xHxW --maxShapes=input:Xx3xHxW --minShape=input:1x3xHxW --saveEngine= --fp16

注意：maxShapes不能大于模型的原始形状。

3. 为TRT引擎编写Deepstream配置文件。


 
参考：
- https://github.com/eriklindernoren/PyTorch-YOLOv3
- https://github.com/marvis/pytorch-caffe-darknet-convert
- https://github.com/marvis/pytorch-yolo3

@article{yolov4, title={YOLOv4: YOLOv4：目标检测的最佳速度与精度}, author={Alexey Bochkovskiy, Chien-Yao Wang, Hong-Yuan Mark Liao}, journal = {arXiv}, year={2020} } ```

PyTorch-YOLOv4 快速上手指南

本指南基于 pytorch-YOLOv4 开源项目，帮助开发者快速在 PyTorch 环境中部署、训练及推理 YOLOv4 模型。

1. 环境准备

系统要求

• Python: 3.6 或 3.7
• PyTorch: 1.4.0 (推荐用于 TensorRT 7.0+)，也支持 1.5.0/1.6.0
• GPU: 支持 CUDA 的 NVIDIA 显卡（训练和加速推理必需）

前置依赖

请确保已安装以下基础库：

pip install torch torchvision opencv-python numpy

若需进行模型转换（ONNX/TensorRT），还需安装：

pip install onnx onnxruntime

2. 安装步骤

2.1 获取代码

克隆项目仓库：

git clone https://github.com/Tianxiaomo/pytorch-YOLOv4.git
cd pytorch-YOLOv4

2.2 下载预训练权重

你可以选择手动转换 Darknet 权重，或直接下载已转换好的 PyTorch 权重文件。

推荐：直接下载已转换的权重（国内用户推荐使用百度网盘）

• 完整模型 (yolov4.pth)

  • 百度网盘: 下载链接 (提取码: xrq9)
  • Google Drive: 下载链接

• 卷积骨干网络 (yolov4.conv.137.pth，用于微调训练)

  • 百度网盘: 下载链接 (提取码: kcel)
  • Google Drive: 下载链接

将下载的 .pth 文件放入项目根目录或 weight 文件夹中。

3. 基本使用

3.1 快速推理 (Inference)

使用预训练权重对单张图片进行检测是最简单的上手方式。

方式一：加载 PyTorch 权重 (.pth) 适用于直接使用本项目提供的权重文件。

python models.py 80 weight/yolov4.pth dog.jpg 416 416

参数说明：80 为类别数 (COCO)，dog.jpg 为输入图片，416 416 为输入宽高。

方式二：加载 Darknet 配置和权重 (.cfg + .weights) 适用于使用官方 Darknet 原始权重文件。

python demo.py -cfgfile cfg/yolov4.cfg -weightfile weight/yolov4.weights -imgfile dog.jpg

3.2 训练自己的数据 (Train)

若需训练自定义数据集，请按以下步骤操作：

1. 准备数据： 生成 train.txt，格式如下（图片路径后跟随边界框坐标和类别 ID）：

   image_path1 x1,y1,x2,y2,id x1,y1,x2,y2,id ...
   image_path2 x1,y1,x2,y2,id x1,y1,x2,y2,id ...
   

   注：可使用 tool/coco_annotation.py 自动转换 COCO 格式数据集。

2. 配置参数： 编辑 cfg.py 文件，设置数据集路径、类别数、batch size 等参数。

3. 开始训练：

   python train.py -g 0 -dir /path/to/your/dataset
   

   参数 -g 指定 GPU ID，-dir 指定数据目录。

3.3 模型导出 (可选)

本项目支持将模型转换为 ONNX 或 TensorRT 引擎以加速推理。

• 转为 ONNX:

  python demo_pytorch2onnx.py weight/yolov4.pth dog.jpg 1 80 416 416
  

• 转为 TensorRT (需安装 TensorRT):

  trtexec --onnx=yolov4.onnx --explicitBatch --saveEngine=yolov4.engine --fp16