mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-26 22:35:43 +00:00
make optimized_model_path be in temp folder instead of source model folder for transformer optimization (#16531)
### The optimize_model will generate a temporary model in current model folder. Most of time, it is fine. However, the scenario will break when the function run against input model mount from AzureML. In that case, the mounted folder is read-only. We have to copy the model to another temp folder to call optimize_model to workaround this issue. Otherwise, the optimize_model will fail when creating the optimized model in the read-only folder. However, the model copy is painful, especially when model is huge. This PR just expose the optimized_model_path at optimize_model level so that the caller could decide where to save the temp model.
This commit is contained in:
parent
2fd25de360
commit
aeaa1d650f
1 changed files with 8 additions and 4 deletions
|
|
@ -20,6 +20,7 @@
|
|||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Dict, Optional
|
||||
|
||||
import coloredlogs
|
||||
|
|
@ -252,6 +253,9 @@ def optimize_model(
|
|||
# stable.
|
||||
disabled_optimizers = ["ConstantSharing"]
|
||||
temp_model_path = None
|
||||
temp_dir = tempfile.TemporaryDirectory()
|
||||
optimized_model_name = "model_o{}_{}.onnx".format(opt_level, "gpu" if use_gpu else "cpu")
|
||||
optimized_model_path = os.path.join(temp_dir.name, optimized_model_name)
|
||||
if opt_level > 1:
|
||||
# Disable some optimizers that might cause failure in symbolic shape inference or attention fusion.
|
||||
disabled_optimizers += (
|
||||
|
|
@ -271,6 +275,7 @@ def optimize_model(
|
|||
opt_level=opt_level,
|
||||
disabled_optimizers=disabled_optimizers,
|
||||
verbose=verbose,
|
||||
optimized_model_path=optimized_model_path,
|
||||
)
|
||||
elif opt_level == 1:
|
||||
# basic optimizations (like constant folding and cast elimination) are not specified to execution provider.
|
||||
|
|
@ -281,6 +286,7 @@ def optimize_model(
|
|||
opt_level=1,
|
||||
disabled_optimizers=disabled_optimizers,
|
||||
verbose=verbose,
|
||||
optimized_model_path=optimized_model_path,
|
||||
)
|
||||
|
||||
if only_onnxruntime and not temp_model_path:
|
||||
|
|
@ -293,10 +299,8 @@ def optimize_model(
|
|||
else:
|
||||
optimizer = optimize_by_fusion(model, model_type, num_heads, hidden_size, optimization_options)
|
||||
|
||||
# Remove the temporary model.
|
||||
if temp_model_path:
|
||||
os.remove(temp_model_path)
|
||||
logger.debug(f"Remove temporary model: {temp_model_path}")
|
||||
# remove the temporary directory
|
||||
temp_dir.cleanup()
|
||||
|
||||
return optimizer
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue