Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.

43 Zeilen
1.6KB

  1. """Exports a pytorch *.pt model to *.onnx format
  2. Usage:
  3. import torch
  4. $ export PYTHONPATH="$PWD" && python models/onnx_export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
  5. """
  6. import argparse
  7. import onnx
  8. from models.common import *
  9. if __name__ == '__main__':
  10. parser = argparse.ArgumentParser()
  11. parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')
  12. parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')
  13. parser.add_argument('--batch-size', type=int, default=1, help='batch size')
  14. opt = parser.parse_args()
  15. print(opt)
  16. # Parameters
  17. f = opt.weights.replace('.pt', '.onnx') # onnx filename
  18. img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size, (1, 3, 320, 192) iDetection
  19. # Load pytorch model
  20. google_utils.attempt_download(opt.weights)
  21. model = torch.load(opt.weights, map_location=torch.device('cpu'))['model']
  22. model.eval()
  23. model.fuse()
  24. # Export to onnx
  25. model.model[-1].export = True # set Detect() layer export=True
  26. _ = model(img) # dry run
  27. torch.onnx.export(model, img, f, verbose=False, opset_version=11, input_names=['images'],
  28. output_names=['output']) # output_names=['classes', 'boxes']
  29. # Check onnx model
  30. model = onnx.load(f) # load onnx model
  31. onnx.checker.check_model(model) # check onnx model
  32. print(onnx.helper.printable_graph(model.graph)) # print a human readable representation of the graph
  33. print('Export complete. ONNX model saved to %s\nView with https://github.com/lutzroeder/netron' % f)