Как преобразовать pytorch модель в onnx?

Question

[Марат Нагаев]

Код:

import torch
import torchvision

dummy_input = torch.randn(1, 3, 416, 416, device="cpu")

model = torch.hub.load('ultralytics/yolov5', 'custom', path='input.pt', force_reload=True) 

# only set that subset of names, starting from the beginning.
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]
torch.onnx.export(model, dummy_input, "output_convert2.onnx", verbose=True, 
opset_version=11, input_names=input_names, output_names=output_names, do_constant_folding=True)

Выдает:
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!