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'.
您好,我在自己学习相关的知识,然后自己修改了一下代码,但很可惜一直报错,我把detect_plate.py中的main修改如下:
不过报错如下:
可以指点下吗?