0
点赞
收藏
分享

微信扫一扫

YOLOX导出onnx文件


导出onnx文件

第一步先确认类别,在yolox/exp/yolox_base.py下面,导出模型的类别数和这里的num_classes保持一致。

YOLOX导出onnx文件_python

本例使用yolo_s模型,yolox自带导出onnx的文件,直接执行命令:

python tools/export_onnx.py --output-name yolox_s.onnx -f  exps/default/yolox_s.py -c YOLOX_outputs/yolox_voc_s/best_ckpt.pth

export_onnx.py

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import argparse
import os
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


def make_parser():
parser = argparse.ArgumentParser("YOLOX onnx deploy")
parser.add_argument(
"--output-name", type=str, default="yolox.onnx", help="output name of models"
)
parser.add_argument(
"--input", default="images", type=str, help="input node name of onnx model"
)
parser.add_argument(
"--output", default="output", type=str, help="output node name of onnx model"
)
parser.add_argument(
"-o", "--opset", default=11, type=int, help="onnx opset version"
)
parser.add_argument("--batch-size", type=int, default=1, help="batch size")
parser.add_argument(
"--dynamic", action="store_true", help="whether the input shape should be dynamic or not"
)
parser.add_argument("--no-onnxsim", action="store_true", help="use onnxsim or not")
parser.add_argument(
"-f",
"--exp_file",
default=None,
type=str,
help="expriment 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,
)

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 = False

logger.info("loading checkpoint done.")
dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])

torch.onnx._export(
model,
dummy_input,
args.output_name,
input_names=[args.input],
output_names=[args.output],
dynamic_axes={args.input: {0: 'batch'},
args.output: {0: 'batch'}} if args.dynamic else None,
opset_version=args.opset,
)
logger.info("generated onnx model named {}".format(args.output_name))

if not args.no_onnxsim:
import onnx

from onnxsim import simplify

input_shapes = {args.input: list(dummy_input.shape)} if args.dynamic else None

# use onnxsimplify to reduce reduent model.
onnx_model = onnx.load(args.output_name)
model_simp, check = simplify(onnx_model,
dynamic_input_shape=args.dynamic,
input_shapes=input_shapes)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, args.output_name)
logger.info("generated simplified onnx model named {}".format(args.output_name))


if __name__ == "__main__":
main()



举报

相关推荐

0 条评论