Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add ncnn deployment examples #145

Merged
merged 14 commits into from
Jul 30, 2021
25 changes: 25 additions & 0 deletions deployment/ncnn/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
cmake_minimum_required(VERSION 3.14)

project(yolort_ncnn)

find_package(OpenCV REQUIRED)

# If the package has been found, several variables will
# be set, you can find the full list with descriptions
# in the OpenCVConfig.cmake file.
# Print some message showing some of them
message(STATUS "OpenCV library status:")
message(STATUS " config: ${OpenCV_DIR}")
message(STATUS " version: ${OpenCV_VERSION}")
message(STATUS " libraries: ${OpenCV_LIBS}")
message(STATUS " include path: ${OpenCV_INCLUDE_DIRS}")

find_package(ncnn REQUIRED)

FILE(GLOB YOLO_SOURCE_FILES *.cpp)

add_executable(yolort_ncnn ${YOLO_SOURCE_FILES})

target_compile_features(yolort_ncnn PUBLIC cxx_range_for)

target_link_libraries(yolort_ncnn ncnn ${OpenCV_LIBS})
48 changes: 48 additions & 0 deletions deployment/ncnn/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
# Ncnn Inference

The ncnn inference for `yolort`, both GPU and CPU are supported.

## Dependencies

- Ubuntu 18.04
- ncnn
- OpenCV 3.4+

## Usage

1. First, Setup the environment variables.

```bash
export TORCH_PATH=$(dirname $(python -c "import torch; print(torch.__file__)"))
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TORCH_PATH/lib/
```

1. First, compile `ncnn` using the following scripts.

```bash
git clone --recursive [email protected]:Tencent/ncnn.git
cd ncnn
mkdir build && cd build
cmake -DCMAKE_BUILD_TYPE=Release -DNCNN_SYSTEM_GLSLANG=ON -DNCNN_BUILD_EXAMPLES=ON .. # Set -DNCNN_VULKAN=ON if you're using VULKAN
make -j4
make install
```

Or follow the [official instructions](https://github.com/Tencent/ncnn/wiki/how-to-build) to install ncnn.

1. Then compile the source code.

```bash
cd deployment/ncnn
mkdir build && cd build
cmake .. -Dncnn_DIR=<ncnn_install_dir>/lib/cmake/ncnn/
make
```

_Note: you have to change <ncnn_install_dir> to your machine's directory, it is the directory that contains ncnnConfig.cmake, if you are following the above operations, you should set it to <./ncnn/build/install>_

1. Now, you can infer your own images with ncnn.

```bash
./yolort_ncnn ../../../test/assets/zidane.jpg
```
62 changes: 62 additions & 0 deletions deployment/ncnn/export_onnx.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
import argparse
import torch
from tools.yolort_deploy_friendly import yolov5s_r40_deploy_ncnn


def get_parser():
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./yolov5s.pt',
help='weights path')
parser.add_argument('--output_path', type=str, default='./yolov5s.onnx',
help='path of exported onnx')
parser.add_argument('--img_size', nargs='+', type=int, default=[640, 640],
help='image (height, width)')
parser.add_argument('--num_classes', type=int, default=80,
help='number of classes')
parser.add_argument('--batch_size', type=int, default=1,
help='batch size')
parser.add_argument('--device', default='cpu',
help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--half', action='store_true',
help='FP16 half-precision export')
parser.add_argument('--dynamic', action='store_true',
help='ONNX: dynamic axes')
parser.add_argument('--simplify', action='store_true',
help='ONNX: simplify model')
parser.add_argument('--opset', type=int, default=11,
help='ONNX: opset version')
return parser


def cli_main():
parser = get_parser()
args = parser.parse_args()
print(args)
export_onnx(args)


def export_onnx(args):

model = yolov5s_r40_deploy_ncnn(
pretrained=True,
num_classes=args.num_classes,
)
img = torch.rand(args.batch_size, 3, 640, 640)
outputs = model(img)
assert len(outputs) == 3

torch.onnx.export(
model,
img,
args.output_path,
verbose=False,
opset_version=args.opset,
do_constant_folding=True,
input_names=['images'],
output_names=['h1', 'h2', 'h3'],
)


if __name__ == "__main__":
cli_main()
Loading