Add executorch support

README.md

@@ -215,6 +215,7 @@ python -m yolox.tools.eval -n  yolox-s -c yolox_s.pth -b 1 -d 1 --conf 0.001 --f
 4. [ncnn in C++ and Java](./demo/ncnn)
 5. [OpenVINO in C++ and Python](./demo/OpenVINO)
 6. [Accelerate YOLOX inference with nebullvm in Python](./demo/nebullvm)
+7. [Executorch export and an Executorch Runtime](./demo/executorch)
 
 ## Third-party resources
 * YOLOX for streaming perception: [StreamYOLO (CVPR 2022 Oral)](https://github.com/yancie-yjr/StreamYOLO)

demo/executorch/README.md

@@ -0,0 +1,60 @@
+## YOLOX-executorch in Python
+
+This doc introduces how to convert your pytorch model into executorch, and how to run an executorch runtime demo to verify your convertion.
+
+### Step1: Install executorch
+
+run the following command to install onnxruntime:
+```shell
+pip install executorch
+```
+
+#### Convert Your Model to ONNX
+
+First, you should move to <YOLOX_HOME> by:
+```shell
+cd <YOLOX_HOME>
+```
+Then, you can:
+
+1. Convert a standard YOLOX model by -n:
+```shell
+python3 tools/export_executorch.py --output-name yolox_s.pte -n yolox-s -c yolox_s.pth
+```
+Notes:
+* -n: specify a model name. The model name must be one of the [yolox-s,m,l,x and yolox-nano, yolox-tiny, yolov3]
+* -c: the model you have trained
+* To customize an input shape for onnx model,  modify the following code in tools/export_executorch.py:
+
+    ```python
+    dummy_input = torch.randn(1, 3, exp.test_size[0], exp.test_size[1])
+    ```
+
+1. Convert a standard YOLOX model by -f. When using -f, the above command is equivalent to:
+
+```shell
+python3 tools/export_executorch.py --output-name yolox_s.pte -f exps/default/yolox_s.py -c yolox_s.pth
+```
+
+3. To convert your customized model, please use -f:
+
+```shell
+python3 tools/export_executorch.py --output-name your_yolox.pte -f exps/your_dir/your_yolox.py -c your_yolox.pth
+```
+
+### Step3: Executorch Runtime Demo
+
+Step1.
+```shell
+cd <YOLOX_HOME>/demo/executorch
+```
+
+Step2. 
+```shell
+python3 executorch_inference.py -m <EXECUTORCH_MODEL_PATH> -i <IMAGE_PATH> -o <OUTPUT_DIR> -s 0.3 --input_shape 640,640
+```
+Notes:
+* -m: your converted pte model
+* -i: input_image
+* -s: score threshold for visualization.
+* --input_shape: should be consistent with the shape you used for executorch convertion.

demo/executorch/executorch_inference.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import argparse
+import os
+
+import cv2
+import numpy as np
+import torch
+
+from executorch.runtime import Runtime
+
+from yolox.data.data_augment import preproc as preprocess
+from yolox.data.datasets import COCO_CLASSES
+from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis
+
+
+def make_parser():
+    parser = argparse.ArgumentParser("executorch inference sample")
+    parser.add_argument(
+        "-m",
+        "--model",
+        type=str,
+        default="yolox.pte",
+        help="Input your pte model.",
+    )
+    parser.add_argument(
+        "-i",
+        "--image_path",
+        type=str,
+        default='test_image.png',
+        help="Path to your input image.",
+    )
+    parser.add_argument(
+        "-o",
+        "--output_dir",
+        type=str,
+        default='demo_output',
+        help="Path to your output directory.",
+    )
+    parser.add_argument(
+        "-s",
+        "--score_thr",
+        type=float,
+        default=0.3,
+        help="Score threshould to filter the result.",
+    )
+    parser.add_argument(
+        "--input_shape",
+        type=str,
+        default="640,640",
+        help="Specify an input shape for inference.",
+    )
+    return parser
+
+
+if __name__ == '__main__':
+    args = make_parser().parse_args()
+
+    input_shape = tuple(map(int, args.input_shape.split(',')))
+    origin_img = cv2.imread(args.image_path)
+    img, ratio = preprocess(origin_img, input_shape)
+
+    runtime = Runtime.get()
+    method = runtime.load_program(args.model).load_method("forward")
+
+    output = method.execute([torch.from_numpy(img).unsqueeze(0)])
+    output = [o.numpy() for o in output]
+
+    predictions = demo_postprocess(output[0], input_shape)[0]
+
+    boxes = predictions[:, :4]
+    scores = predictions[:, 4:5] * predictions[:, 5:]
+
+    boxes_xyxy = np.ones_like(boxes)
+    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
+    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
+    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
+    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
+    boxes_xyxy /= ratio
+    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
+    if dets is not None:
+        final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
+        origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
+                         conf=args.score_thr, class_names=COCO_CLASSES)
+
+    mkdir(args.output_dir)
+    output_path = os.path.join(args.output_dir, os.path.basename(args.image_path))
+    cv2.imwrite(output_path, origin_img)

tools/export_executorch.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import argparse
+import os
+from xml.parsers.expat import model
+from loguru import logger
+
+import torch
+from torch import nn
+
+from yolox.exp import get_exp
+from yolox.models.network_blocks import SiLU
+from yolox.utils import replace_module
+
+
+from executorch.exir import to_edge_transform_and_lower
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
+
+
+def make_parser():
+    parser = argparse.ArgumentParser("YOLOX executorch deploy")
+    parser.add_argument(
+        "--output-name", type=str, default="yolox.pte", help="output name of models"
+    )
+    parser.add_argument("--batch-size", type=int, default=1, help="batch size")
+    parser.add_argument(
+        "-f",
+        "--exp_file",
+        default=None,
+        type=str,
+        help="experiment description file",
+    )
+    parser.add_argument("-expn", "--experiment-name", type=str, default=None)
+    parser.add_argument("-n", "--name", type=str, default=None, help="model name")
+    parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path")
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    parser.add_argument(
+        "--decode_in_inference",
+        action="store_true",
+        help="decode in inference or not"
+    )
+
+    return parser
+
+
+@logger.catch
+def main():
+    args = make_parser().parse_args()
+    logger.info("args value: {}".format(args))
+    exp = get_exp(args.exp_file, args.name)
+    exp.merge(args.opts)
+
+    if not args.experiment_name:
+        args.experiment_name = exp.exp_name
+
+    model = exp.get_model()
+    if args.ckpt is None:
+        file_name = os.path.join(exp.output_dir, args.experiment_name)
+        ckpt_file = os.path.join(file_name, "best_ckpt.pth")
+    else:
+        ckpt_file = args.ckpt
+
+    # load the model state dict
+    ckpt = torch.load(ckpt_file, map_location="cpu")
+
+    model.eval()
+    if "model" in ckpt:
+        ckpt = ckpt["model"]
+    model.load_state_dict(ckpt)
+    model = replace_module(model, nn.SiLU, SiLU)
+    model.head.decode_in_inference = args.decode_in_inference
+
+    logger.info("loading checkpoint done.")
+    dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])
+
+    exported_program = torch.export.export(model, (dummy_input,))
+
+    program = to_edge_transform_and_lower(
+        exported_program,
+        partitioner=[XnnpackPartitioner()]  # CPU | CoreMLPartitioner() for iOS | QnnPartitioner() for Qualcomm
+    ).to_executorch()
+
+    with open(args.output_name, "wb") as f:
+        f.write(program.buffer)
+
+    logger.info("generated executorch model named {}".format(args.output_name))
+
+if __name__ == "__main__":
+    main()