saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support export >2G model when using optimizer.py only (#6014)

Browse source 
* checkin

* add warning if user specify same inut and output path
This commit is contained in:
Ye Wang 2020-12-07 17:18:49 -08:00 committed by GitHub
parent b348538c8a
commit fa06be2133
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
onnxruntime/python/tools/transformers/fusion_reshape.py
View file

@ -126,7 +126,8 @@ class FusionReshape(Fusion):
value=helper.make_tensor(name='const_tensor',
data_type=TensorProto.INT64,
dims=shape_value.shape,
vals=shape_value))
vals=bytes(shape_value),
raw=True))
reshape_node.input[1] = constant_shape_name
reshape_node.name = self.model.create_node_name('Reshape', 'Reshape_Fuse')
self.nodes_to_remove.extend([concat_node])

10
onnxruntime/python/tools/transformers/optimizer.py
View file

@ -214,6 +214,9 @@ def _parse_arguments():
default=0,
help="onnxruntime optimization level. 0 will disable onnxruntime.")
parser.add_argument('--use_external_data_format', required=False, action='store_true', help="use external data format")
parser.set_defaults(use_external_data_format=False)
args = parser.parse_args()
return args
@ -321,6 +324,11 @@ def main():
_setup_logger(args.verbose)
if os.path.realpath(args.input) == os.path.realpath(args.output):
logger.warning(
f"Specified the same input and output path. Note that this may overwrite the original model"
)
optimization_options = _get_optimization_options(args)
optimizer = optimize_model(args.input,
@ -338,7 +346,7 @@ def main():
if args.input_int32:
optimizer.change_input_to_int32()
optimizer.save_model_to_file(args.output)
optimizer.save_model_to_file(args.output, args.use_external_data_format)
if optimizer.is_fully_optimized():
logger.info("The model has been fully optimized.")