zeusees / License-Plate-Detector

基于Yolov5车牌检测,更快更准.
http://www.zeusee.com
1.21k stars 400 forks source link

请教如何转为onnx #36

Closed moneypi closed 2 years ago

moneypi commented 2 years ago

您好,我在自己学习相关的知识,然后自己修改了一下代码,但很可惜一直报错,我把detect_plate.py中的main修改如下:

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp/weights/last.pt', help='model.pt path(s)')
    parser.add_argument('--image', type=str, default='data/images/test.jpg', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    opt = parser.parse_args()
    print(opt)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # print(opt.weights)
    model = torch.load(opt.weights[0], map_location=device)['model'].float().fuse().eval()
    # model = torch.load(opt.weights)

    batch_size = 1                              # 批处理大小
    input_shape = (3, 800, 800)                 # 输入数据

    # set the model to inference mode
    # torch_model.eval()

    x = torch.randn(batch_size,*input_shape)    # 生成张量
    export_onnx_file = "test.onnx"              # 目的ONNX文件名
    torch.onnx.export(model,
                        x,
                        export_onnx_file,
                        opset_version=10,
                        do_constant_folding=True,   # 是否执行常量折叠优化
                        input_names=["input"],      # 输入名
                        output_names=["output"])

    quit()
    for filename in f:
        print(filename)
        detect_one(model, filename, device)

不过报错如下:

Namespace(image='data/images/test.jpg', img_size=640, weights=['weights/plate_det_model.pt'])
Fusing layers... 
[W NNPACK.cpp:80] Could not initialize NNPACK! Reason: Unsupported hardware.
/home/luke/Downloads/License-Plate-Detector/models/yolo.py:62: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if self.grid[i].shape[2:4] != x[i].shape[2:4]:
Traceback (most recent call last):
  File "detect_plate.py", line 185, in <module>
    torch.onnx.export(model,
  File "/home/luke/miniconda3/lib/python3.8/site-packages/torch/onnx/__init__.py", line 225, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "/home/luke/miniconda3/lib/python3.8/site-packages/torch/onnx/utils.py", line 85, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "/home/luke/miniconda3/lib/python3.8/site-packages/torch/onnx/utils.py", line 632, in _export
    _model_to_graph(model, args, verbose, input_names,
  File "/home/luke/miniconda3/lib/python3.8/site-packages/torch/onnx/utils.py", line 417, in _model_to_graph
    graph = _optimize_graph(graph, operator_export_type,
  File "/home/luke/miniconda3/lib/python3.8/site-packages/torch/onnx/utils.py", line 168, in _optimize_graph
    torch._C._jit_pass_onnx_prepare_inplace_ops_for_onnx(graph)
RuntimeError: 

aten::view(Tensor(a) self, int[] size) -> (Tensor(a)):
Expected a value of type 'Tensor' for argument 'self' but instead found type 'None'.

可以指点下吗?

moneypi commented 2 years ago

基本环境为: Python:3.8.10 Pytorch:1.7.1 Ubuntu: 21.10

Jinkham commented 2 years ago

@moneypi models文件夹下有对应的转化代码