mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

export.py update

pull/275/head
Glenn Jocher 2020-07-01 16:14:49 -07:00
parent f1d67f4110
commit a62a1c2c67
2 changed files with 9 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
detect.py
View File

@ -156,9 +156,9 @@ if __name__ == '__main__':
print(opt) print(opt)
with torch.no_grad(): with torch.no_grad():
detect() # detect()
# Update all models # Update all models
# for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov3-spp.pt']: for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov3-spp.pt']:
# detect() detect()
# create_pretrained(opt.weights, opt.weights) create_pretrained(opt.weights, opt.weights)

10
models/export.py
View File

@ -6,8 +6,6 @@ Usage:
import argparse import argparse
import onnx
from models.common import * from models.common import *
from utils import google_utils from utils import google_utils
@ -21,7 +19,7 @@ if __name__ == '__main__':
print(opt) print(opt)
# Input # Input
img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size, (1, 3, 320, 192) iDetection img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection
# Load PyTorch model # Load PyTorch model
google_utils.attempt_download(opt.weights) google_utils.attempt_download(opt.weights)
@ -30,7 +28,7 @@ if __name__ == '__main__':
model.model[-1].export = True # set Detect() layer export=True model.model[-1].export = True # set Detect() layer export=True
_ = model(img) # dry run _ = model(img) # dry run
# Export to TorchScript # TorchScript export
try: try:
f = opt.weights.replace('.pt', '.torchscript') # filename f = opt.weights.replace('.pt', '.torchscript') # filename
ts = torch.jit.trace(model, img) ts = torch.jit.trace(model, img)
@ -39,8 +37,10 @@ if __name__ == '__main__':
except Exception as e: except Exception as e:
print('TorchScript export failed: %s' % e) print('TorchScript export failed: %s' % e)
# Export to ONNX # ONNX export
try: try:
import onnx
f = opt.weights.replace('.pt', '.onnx') # filename f = opt.weights.replace('.pt', '.onnx') # filename
model.fuse() # only for ONNX model.fuse() # only for ONNX
torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],