From 982b5027b992cbeb8a5a4bb001817c4038690bca Mon Sep 17 00:00:00 2001 From: Ibai Date: Fri, 8 Apr 2022 21:52:13 +0900 Subject: [PATCH] Fix flow_width --- convert_to_onnx.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/convert_to_onnx.py b/convert_to_onnx.py index 79e06ac..dce0737 100644 --- a/convert_to_onnx.py +++ b/convert_to_onnx.py @@ -20,18 +20,7 @@ if __name__ == '__main__': t1 = torch.rand(1, 3, in_h, in_w) t2 = torch.rand(1, 3, in_h, in_w) - flow_init = torch.rand(1, 2, in_h//2, in_w) - - # Export the model - # !! Needs Pytorch nightly until next release (1.12). Ref: https://github.com/pytorch/pytorch/pull/73760 - torch.onnx.export(model, - (t1_half, t2_half), - "crestereo_without_flow.onnx", # where to save the model (can be a file or file-like object) - export_params=True, # store the trained parameter weights inside the model file - opset_version=12, # the ONNX version to export the model to - do_constant_folding=True, # whether to execute constant folding for optimization - input_names = ['left', 'right'], # the model's input names - output_names = ['output']) + flow_init = torch.rand(1, 2, in_h//2, in_w//2) # Export the model torch.onnx.export(model, @@ -43,5 +32,17 @@ if __name__ == '__main__': input_names = ['left', 'right','flow_init'], # the model's input names output_names = ['output']) + # # Export the model without init_flow (it takes a lot of time) + # # !! Needs Pytorch nightly until next release (1.12). Ref: https://github.com/pytorch/pytorch/pull/73760 + # torch.onnx.export(model, + # (t1_half, t2_half), + # "crestereo_without_flow.onnx", # where to save the model (can be a file or file-like object) + # export_params=True, # store the trained parameter weights inside the model file + # opset_version=12, # the ONNX version to export the model to + # do_constant_folding=True, # whether to execute constant folding for optimization + # input_names = ['left', 'right'], # the model's input names + # output_names = ['output']) + +