Skip to content

Refactor resnet to use @capture_outputs / @can_return_tuple output tracing#44019

Open
Sid-V5 wants to merge 1 commit intohuggingface:mainfrom
Sid-V5:refactor-resnet-output-tracing
Open

Refactor resnet to use @capture_outputs / @can_return_tuple output tracing#44019
Sid-V5 wants to merge 1 commit intohuggingface:mainfrom
Sid-V5:refactor-resnet-output-tracing

Conversation

@Sid-V5
Copy link
Copy Markdown

@Sid-V5 Sid-V5 commented Feb 15, 2026

Refactored the resnet model to use the standardized output tracing decorators (@capture_outputs and @can_return_tuple) as part of the migration

Changes

File Change
modeling_resnet.py Migrated to @capture_outputs / @can_return_tuple

Specifically:

  1. Added _can_record_outputs = {"hidden_states": ResNetStage} to ResNetPreTrainedModel hooks on ResNetStage submodules automatically collect hidden states.
  2. Decorated ResNetModel.forward() with @capture_outputs removes manual output_hidden_states / return_dict parameter resolution and collection loops.
  3. Decorated ResNetForImageClassification.forward() and ResNetBackbone.forward() with @can_return_tuple handles return_dict conversion without hook installation.
  4. Simplified ResNetEncoder.forward() now returns a single hidden_state tensor instead of manually building output tuples/dicts.
  5. Removed explicit output_hidden_states, return_dict parameters from internal signatures where they are no longer needed.
  6. Added imports for can_return_tuple and capture_outputs.

Reference implementations

Follows the pattern established in llama, mistral, and qwen2 models as described in #43979.

Notes

  • ResNet has no attention layers, so _can_record_outputs only maps "hidden_states" (no "attentions" entry needed).
  • The ResNetBackbone class uses BackboneMixin which has its own output handling the @can_return_tuple decorator is applied for consistency.

Fixes (resnet portion)

Replace manual `output_hidden_states` / `return_dict` boilerplate with
`@capture_outputs` on the base model and `@can_return_tuple` on wrapper
forwards, following the pattern established in llama/mistral/qwen2.

- Add `_can_record_outputs = {"hidden_states": ResNetStage}` to
  `ResNetPreTrainedModel`
- Simplify `ResNetEncoder.forward()` to return a single tensor
- Decorate `ResNetModel.forward()` with `@capture_outputs`
- Decorate `ResNetForImageClassification.forward()` and
  `ResNetBackbone.forward()` with `@can_return_tuple`
- Remove explicit `output_hidden_states`, `return_dict` parameter
  resolution and manual collection loops

Closes huggingface#43979 (resnet portion)
@github-actions
Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: resnet

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant