Skip to content

Commit

Permalink
[rknn] update rkopt doc and code, now model is same for tk1/2
Browse files Browse the repository at this point in the history
  • Loading branch information
zen committed Jul 25, 2023
1 parent 4390b71 commit ff5ef69
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 13 deletions.
6 changes: 6 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,9 @@
#### Get model optimized for RKNN
Exports model with optimization for RKNN, please refer here [README_rkopt_manual.md](./README_rkopt_manual.md)

---
<br>

<div align="center">
<p>
<a align="left" href="https://ultralytics.com/yolov5" target="_blank">
Expand Down
39 changes: 29 additions & 10 deletions README_rkopt_manual.md
Original file line number Diff line number Diff line change
@@ -1,16 +1,35 @@
# YOLOv5 - rkopt 仓库
# YOLOv5 - RKNN optimize

- 基于 https://github.com/ultralytics/yolov5 代码修改,设配 rknpu 设备的部署优化
- 切换分支 git checkout {分支名}
- 目前支持分支:
- master
- maxpool/ focus 优化,输出改为个branch分支的输出。以上优化代码使用插入宏实现,不影响原来的训练逻辑,这个优化兼容修改前的权重,故支持官方给的预训练权重。
## Source

- 修改激活函数 silu 为 relu
Base on https://github.com/ultralytics/yolov5 with commit id as d3ea0df8b9f923685ce5f2555c303b8eddbf83fd

- 训练的相关内容请参考 README.md 说明

- 导出模型时 python export.py --rknpu {rk_platform} 即可导出优化模型

(rk_platform支持 rk1808, rv1109, rv1126, rk3399pro, rk3566, rk3568, rk3588, rv1103, rv1106)
## What different

Inference result unchanged:

- Optimize focus/SPPF block, getting better performance with same result
- Change output node, remove post_process from the model. (post process is unfriendly in quantization)



Inference result changed:

- Using ReLU as activation layer instead of SiLU(Only valid when training new model)



## How to use

```
python export.py --rknpu {rk_platform} --weight yolov5s.pt
```

- rk_platform support rk1808, rv1109, rv1126, rk3399pro, rk3566, rk3562, rk3568, rk3588, rv1103, rv1106. (Actually the exported models are the same in spite of the exact platform )

- the 'yolov5s.pt' could be replace with your model path
- A file name "RK_anchors.txt" would be generated and it could be use during doing post_process in the outside.
- **NOTICE: Please call with --rknpu param, do not changing the default rknpu value in export.py.**

2 changes: 1 addition & 1 deletion export.py
Original file line number Diff line number Diff line change
Expand Up @@ -526,7 +526,7 @@ def run(
m.onnx_dynamic = dynamic
m.export = True

if os.getenv('RKNN_model_hack', '0') == 'npu_1':
if os.getenv('RKNN_model_hack', '0') in ['npu_1','npu_2']:
from models.common import Focus
from models.common import Conv
from models.common_rk_plug_in import surrogate_focus
Expand Down
4 changes: 2 additions & 2 deletions models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -232,7 +232,7 @@ def forward(self, x):
return self.cv2(torch.cat(y, 1))


if os.getenv('RKNN_model_hack', '0') in ['0','npu_2']:
if os.getenv('RKNN_model_hack', '0') in ['0']:
class SPPF(nn.Module):
# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
Expand All @@ -249,7 +249,7 @@ def forward(self, x):
y1 = self.m(x)
y2 = self.m(y1)
return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
elif os.getenv('RKNN_model_hack', '0') == 'npu_1':
elif os.getenv('RKNN_model_hack', '0') in ['npu_1','npu_2']:
class SPPF(nn.Module):
# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
Expand Down

0 comments on commit ff5ef69

Please sign in to comment.