WongKinYiu / yolor

implementation of paper - You Only Learn One Representation: Unified Network for Multiple Tasks (https://arxiv.org/abs/2105.04206)
GNU General Public License v3.0
1.99k stars 518 forks source link

使用main分支下models/export.py转换yolor_p6.pt失败 #7

Open mowen101 opened 3 years ago

mowen101 commented 3 years ago

代码修改如下:

import argparse from models import * from utils.torch_utils import select_device import torch from utils.google_utils import attempt_download from utils.google_utils import attempt_load

if name == 'main':

parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./yolov4.pt', help='weights path')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
parser.add_argument('--cfg', type=str, help='batch size')
opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
print(opt)

# Input
img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) ### ### iDetection
# Load PyTorch model
attempt_download(opt.weights)
model = Darknet(opt.cfg, opt.img_size)
model.load_state_dict(torch.load(opt.weights)['model'])
model.eval()
y = model(img)  # dry run

try:
    import onnx

    print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
    f = opt.weights.replace('.pt', '.onnx')  # filename
    model.fuse()  # only for ONNX
    torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
                      output_names=['output'])
    # Checks
    onnx_model = onnx.load(f)  # load onnx model
    onnx.checker.check_model(onnx_model)  # check onnx model
    print(onnx.helper.printable_graph(onnx_model.graph))  # print a human readable model
    print('ONNX export success, saved as %s' % f)
except Exception as e:
    print('ONNX export failure: %s' % e)

代码我修改为上面的样式,onnx版本1.9.0, 使用命令为:python models/export.py --weights yolor_p6.pt --img-size 1280 --cfg cfg/yolor_p6.cfg

目前报错ONNX export failure: Exporting the operator silu to ONNX opset version 12 is not supported. Please open a bug to request ONNX export support for the missing operator.

我尝试了其他版本的onnx,也尝试了不同版本的opset_version,10以上都是这个错误,10以下是ONNX export failure: step!=1 is currently not supported

非常感谢,期待您的回复

WongKinYiu commented 3 years ago

有些新的層還沒支援轉換