ONNX export in .train() mode fix (#3362)
This commit is contained in:
1 parent
2435bfe896
commit
ef4d53818d
1 file changed
+2
@@ -97,6 +97,8 @@ if __name__ == '__main__':
|
|||||||
print(f'{prefix} starting export with onnx {onnx.__version__}...')
|
print(f'{prefix} starting export with onnx {onnx.__version__}...')
|
||||||
f = opt.weights.replace('.pt', '.onnx') # filename
|
f = opt.weights.replace('.pt', '.onnx') # filename
|
||||||
torch.onnx.export(model, img, f, verbose=False, opset_version=opt.opset_version, input_names=['images'],
|
torch.onnx.export(model, img, f, verbose=False, opset_version=opt.opset_version, input_names=['images'],
|
||||||
|
training=torch.onnx.TrainingMode.TRAINING if opt.train else torch.onnx.TrainingMode.EVAL,
|
||||||
|
do_constant_folding=not opt.train,
|
||||||
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
|
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
|
||||||
'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
|
'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user