Retain the parameter names in ONNX exporter #17551
Conversation
houseroad
force-pushed
the
onnx_param_name
branch
from
43910a1 to
481aa30
Compare
torch/onnx/utils.py
Outdated
| @@ -44,7 +44,7 @@ def set_training(model, mode): | |||
|
|
|||
| def export(model, args, f, export_params=True, verbose=False, training=False, | |||
| input_names=None, output_names=None, aten=False, export_raw_ir=False, | |||
| operator_export_type=None, opset_version=None): | |||
| operator_export_type=None, opset_version=None, reserve_param_name=False): | |||
There was a problem hiding this comment.
A minor point - could we consider another name for this parameter, such as retain_param_name or export_with_param_name?
houseroad
force-pushed
the
onnx_param_name
branch
from
481aa30 to
d2465c7
Compare
houseroad
changed the title
torch/onnx/utils.py
Outdated
| param_names = list(state_dict.keys()) | ||
| for i in range(len(graph_inputs)): | ||
| if i >= user_input_num: | ||
| graph_inputs[i].setUniqueName( |
There was a problem hiding this comment.
what happens if there's a duplicated name?
There was a problem hiding this comment.
https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/ir.cpp#L670, it will rename the old owner of this name. We should be good. :-)
torch/onnx/utils.py
Outdated
| graph_inputs = list(graph.inputs()) | ||
| user_input_num = len(graph_inputs) - len(state_dict) | ||
| param_names = list(state_dict.keys()) | ||
| for i in range(len(graph_inputs)): |
There was a problem hiding this comment.
nit:
for i, inp in enumerate(graph_inputs):
inp.setUniqueName(...)
| if retain_param_name: | ||
| graph_inputs = list(graph.inputs()) | ||
| user_input_num = len(graph_inputs) - len(state_dict) | ||
| param_names = list(state_dict.keys()) |
There was a problem hiding this comment.
wait, the order of the keys is not stable
|
I will rebase this pr, and have another PR to enable it as default behavior. Don't want to let the massive updated expect files overwhelmed the important changes. I think to be safe, let's at least keep this flag for now. |
houseroad
force-pushed
the
onnx_param_name
branch
from
d2465c7 to
eba0a02
Compare
There was a problem hiding this comment.
@houseroad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
If you plan this flag as a temporary thing - start the name of the argument with underscore (i.e. _preserve_param_names) |
There was a problem hiding this comment.
@houseroad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
Can you please revert all of the changed |
There was a problem hiding this comment.
@houseroad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
There was a problem hiding this comment.
@houseroad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
There was a problem hiding this comment.
ship it
@sampepose - the .expect file changes are expected as we change the names of the parameters
| graph_inputs = list(graph.inputs()) | ||
| user_input_num = len(graph_inputs) - len(state_dict) | ||
| param_names = list(state_dict.keys()) | ||
| for i, inp in enumerate(graph_inputs): |
There was a problem hiding this comment.
or for inp, name in zip(graph_inputs[user_input_num:], param_names):
So, we will keep the names of ONNX initializers the same as the names in PyTorch state dict.
Later, we will make this as the default behavior.