Skip to content

Commit

Permalink
fix fetch op rename_input bug in QAT export model
Browse files Browse the repository at this point in the history
  • Loading branch information
yghstill committed Dec 10, 2021
1 parent a4c0c71 commit 1e84d5d
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions python/paddle/fluid/contrib/slim/quantization/imperative/qat.py
Original file line number Diff line number Diff line change
Expand Up @@ -484,7 +484,7 @@ def save_quantized_model(self, model, path, input_spec=None, **config):
model_filename=model_filename,
params_filename=params_filename))

self._gather_scales(infer_program, scope)
self._gather_scales(infer_program, scope, fetch_targets)

self._set_skip_quant_attr(infer_program)

Expand Down Expand Up @@ -520,10 +520,10 @@ def _is_target_layer(self, layer):

return flag

def _gather_scales(self, program, scope):
def _gather_scales(self, program, scope, fetch_targets):
"""
Get all scales from fake ops, save them into the corresponding ops
and delete all moving_average_abs_max_scale ops.
and delete all moving_average_abs_max_scale ops.
"""

def _gather_input_scale():
Expand Down Expand Up @@ -580,6 +580,11 @@ def _gather_output_scale():

for next_op in next_ops:
next_op._rename_input(out_var_name, in_var_name)
# If next_op is `fetch` and out_var_name in fetch_targets,
# fetch_targets must update to in_var_name when rename input.
for i in range(len(fetch_targets)):
if fetch_targets[i].name == out_var_name:
fetch_targets[i] = block.var(in_var_name)

_gather_input_scale()
_gather_output_scale()
Expand Down

1 comment on commit 1e84d5d

@paddle-bot-old
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Congratulation! Your pull request passed all required CI. You could ask reviewer(s) to approve and merge. 🎉

Please sign in to comment.