Skip to content

only flatten a pytree once#6767

Merged
datumbox merged 1 commit intopytorch:mainfrom
pmeier:flatten-once
Oct 14, 2022
Merged

only flatten a pytree once#6767
datumbox merged 1 commit intopytorch:mainfrom
pmeier:flatten-once

Conversation

@pmeier
Copy link
Copy Markdown
Contributor

@pmeier pmeier commented Oct 14, 2022

Closes #6760 minus the container handling as stated in #6760 (comment).

h, w = image.spatial_size = (24, 32)

params = transform._get_params(image)
params = transform._get_params([image])
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

All changes in this file stem from the fact that _get_params previously handled pytree objects, but now only handles flattened ones.

[
[transforms.Pad(2), transforms.RandomCrop(28)],
[lambda x: 2.0 * x, transforms.Pad(2), transforms.RandomCrop(28)],
[transforms.Pad(2), lambda x: 2.0 * x, transforms.RandomCrop(28)],
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I've accidentally included this while playing with the container flattening. It is useful nevertheless to have the "foreign" transform also in the middle of the pipeline. Let me know if I should revert.

prototype_transform = prototype_transforms.RandomApply(
[
prototype_transforms.Resize(256),
legacy_transforms.CenterCrop(224),
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This was just plain wrong. It never triggered because the input is a single image. Found this while playing with the container flattening.

return key, dct[key]

def _extract_image_or_video(
def _flatten_and_extract_image_or_video(
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I've opted to group the flattening and extraction as well as the unflattening and insertion since that yielded the cleanest results.

Copy link
Copy Markdown
Contributor

@datumbox datumbox left a comment

Choose a reason for hiding this comment

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

LGTM @pmeier. I like the approach. It keeps this detail internal and minimizes the overhead. Let's see how this improved our speed in real world applications. I'll do some training later to confirm.

@datumbox datumbox merged commit e3238e5 into pytorch:main Oct 14, 2022
@pmeier pmeier deleted the flatten-once branch October 14, 2022 08:22
facebook-github-bot pushed a commit that referenced this pull request Oct 17, 2022
Reviewed By: NicolasHug

Differential Revision: D40427480

fbshipit-source-id: 0552b32b56e5292a64060fcddde46feca4137b6a
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Avoid multiple pytree flattenings inside the prototype transforms

3 participants