mirror
/
MODNet
mirror of https://github.com/ZHKKKe/MODNet.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

remove whitespace changes

pull/217/head
Richard Brown 2024-05-06 16:24:23 +02:00
parent 72ad182ea2
commit 67a565bcd1
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
onnx/export_onnx.py
View File

@ -1,7 +1,7 @@
"""
Export ONNX model of MODNet with:
input shape: (batch_size, 3, height, width)
output shape: (batch_size, 1, height, width)
output shape: (batch_size, 1, height, width)
Arguments:
--ckpt-path: path of the checkpoint that will be converted
@ -50,6 +50,6 @@ if __name__ == '__main__':
# export to onnx model
torch.onnx.export(
modnet.module, dummy_input, args.output_path, export_params = True,
input_names = ['input'], output_names = ['output'],
modnet.module, dummy_input, args.output_path, export_params = True,
input_names = ['input'], output_names = ['output'],
dynamic_axes = {'input': {0:'batch_size', 2:'height', 3:'width'}, 'output': {0: 'batch_size', 2: 'height', 3: 'width'}})