@@ -126,7 +126,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:'
126
126
127
127
128
128
@try_export
129
- def export_onnx (model , im , file , opset , train , dynamic , simplify , prefix = colorstr ('ONNX:' )):
129
+ def export_onnx (model , im , file , opset , dynamic , simplify , prefix = colorstr ('ONNX:' )):
130
130
# YOLOv5 ONNX export
131
131
check_requirements ('onnx' )
132
132
import onnx
@@ -140,8 +140,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
140
140
f ,
141
141
verbose = False ,
142
142
opset_version = opset ,
143
- training = torch .onnx .TrainingMode .TRAINING if train else torch .onnx .TrainingMode .EVAL ,
144
- do_constant_folding = not train ,
143
+ do_constant_folding = True ,
145
144
input_names = ['images' ],
146
145
output_names = ['output' ],
147
146
dynamic_axes = {
@@ -459,7 +458,6 @@ def run(
459
458
include = ('torchscript' , 'onnx' ), # include formats
460
459
half = False , # FP16 half-precision export
461
460
inplace = False , # set YOLOv5 Detect() inplace=True
462
- train = False , # model.train() mode
463
461
keras = False , # use Keras
464
462
optimize = False , # TorchScript: optimize for mobile
465
463
int8 = False , # CoreML/TF INT8 quantization
@@ -501,7 +499,7 @@ def run(
501
499
im = torch .zeros (batch_size , 3 , * imgsz ).to (device ) # image size(1,3,320,192) BCHW iDetection
502
500
503
501
# Update model
504
- model .train () if train else model . eval () # training mode = no Detect() layer grid construction
502
+ model .eval ()
505
503
for k , m in model .named_modules ():
506
504
if isinstance (m , Detect ):
507
505
m .inplace = inplace
@@ -524,7 +522,7 @@ def run(
524
522
if engine : # TensorRT required before ONNX
525
523
f [1 ], _ = export_engine (model , im , file , half , dynamic , simplify , workspace , verbose )
526
524
if onnx or xml : # OpenVINO requires ONNX
527
- f [2 ], _ = export_onnx (model , im , file , opset , train , dynamic , simplify )
525
+ f [2 ], _ = export_onnx (model , im , file , opset , dynamic , simplify )
528
526
if xml : # OpenVINO
529
527
f [3 ], _ = export_openvino (file , metadata , half )
530
528
if coreml : # CoreML
@@ -578,7 +576,6 @@ def parse_opt():
578
576
parser .add_argument ('--device' , default = 'cpu' , help = 'cuda device, i.e. 0 or 0,1,2,3 or cpu' )
579
577
parser .add_argument ('--half' , action = 'store_true' , help = 'FP16 half-precision export' )
580
578
parser .add_argument ('--inplace' , action = 'store_true' , help = 'set YOLOv5 Detect() inplace=True' )
581
- parser .add_argument ('--train' , action = 'store_true' , help = 'model.train() mode' )
582
579
parser .add_argument ('--keras' , action = 'store_true' , help = 'TF: use Keras' )
583
580
parser .add_argument ('--optimize' , action = 'store_true' , help = 'TorchScript: optimize for mobile' )
584
581
parser .add_argument ('--int8' , action = 'store_true' , help = 'CoreML/TF INT8 quantization' )
0 commit comments