Home Original page

Closed the perf gap of resnet and enabled refit by cehongwang · Pull Request #3629 · pytorch/TensorRT

Conversation

@cehongwang

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

peri044

Choose a reason for hiding this comment

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

minor comment. LGTM

@cehongwang cehongwang linked an issue

Jun 27, 2025

that may be closed by this pull request

narendasan

engine_weight_map = {}
for engine_weight_name, (sd_weight_name, np_weight_type) in weight_name_map.items():
if sd_weight_name not in state_dict:
if engine_weight_name.split(" ")[-1] in ["SCALE", "SHIFT"]:

Choose a reason for hiding this comment

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

We should abstract this imo. Like if there are any weight types that require constant folding in converter this should be associated with the converter. Then the refit system will just iterate through all these constant fold operations. Ideally the converter can use the same implementation

Choose a reason for hiding this comment

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

Currently BN is the only one. Do you think we should have a constant_fold function and have refit and conversion call that function?

Choose a reason for hiding this comment

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

yeah we should think of some sort of abstraction where the constant folding process is defined once for both refit and conversion

HolyWu

narendasan

narendasan

Choose a reason for hiding this comment

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

LGTM

Labels