Resnet flax by Shubhamai · Pull Request #21472 · huggingface/transformers

Shubhamai · 2023-02-06T13:11:46Z

Flax Implementation of `microsoft/resnet-50`

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

@sanchit-gandhi I guess 😅

Status

Last Updated - Sunday, 12 February 2023

TODO

Blocked on merge of this PR to add support for BatchNorm layers.
Uploading Shubhamai/resnet-50 flax weights to microsoft/resnet-50.

HuggingFaceDocBuilderDev · 2023-02-06T13:26:46Z

The documentation is not available anymore as the PR was closed or merged.

sanchit-gandhi · 2023-02-22T13:45:35Z

Very cool! Sorry to only reply here now - looks like you've made a really solid start to this PR! Let's get the Batch Norm PR merged ASAP and then go full send on Flax ResNet! 🚀

Feel free to ping me with any questions / queries! More than happy to help with the integration!

Shubhamai · 2023-02-24T09:26:32Z

@sanchit-gandhi the PR is also now ready for your review, thanks a lot for your time.

sanchit-gandhi

Hey @shubhamagarwal92! Super nice PR! Think you've handled all the conversions of nn.Sequential -> LayerCollection really well!

Left a few comments regarding i) Flax train/deterministic naming conventions, ii) concerns about control flow and XLA compatibility, and iii) one or two unused args that we can remove.

These are very specific because I think the PR is otherwise in very good shape! Feel free to ask any questions / queries if you have them, more than happy to help with the final bit of integration here!

sanchit-gandhi · 2023-02-27T16:34:17Z

+            padding=self.kernel_size // 2,
+            dtype=self.dtype,
+            use_bias=False,
+            kernel_init=nn.initializers.variance_scaling(2.0, mode="fan_out", distribution="normal", dtype=self.dtype),


Is this the correct kernel_init for ResNet? It might very well be! I'm not a computer vision expert so am not familiar what the standard convolutional inits are here, but what I do know is that for Flax VIT we use a slightly different configuration:

transformers/src/transformers/models/vit/modeling_flax_vit.py

Line 103 in 92dfceb

kernel_init=jax.nn.initializers.variance_scaling(

After changing normal to truncated_normal (in the upcoming commit), I think it's good to go.

In pt implementation of resnet, we used nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu") to initialize Conv2D. Jax does provide partial implementation of kaiming_normal_ as he_normal, partial because it currently does not provide the nonlinearity parameter.

he_normal as mentioned in docs, is simply a specialization of jax.nn.initializers.variance_scaling() . Since pt implementation uses fan_out mode. I had to use variance_scaling initializer for more "parameter control".

normal is changed to truncated_normal because, the same is done in the he_normal implementation of Jax.

I don't remember the last time I used initializers (esp. since normalization and skip connections), so I am also probably not the right person for this, but from the current knowledge I encompass, it looks good to go.

Awesome - thanks for the explanation! Sounds good to me!

sanchit-gandhi · 2023-02-27T17:21:25Z

+
+        for stage_module in self.stages:
+            if output_hidden_states:
+                hidden_states = hidden_states + (hidden_state.transpose(0, 3, 1, 2),)


Think it's ok that we transpose here since we don't do any ops on the hidden state, just return it as part of the dict (is not going to harm XLA)

Shubhamai · 2023-02-28T09:25:09Z

Thanks for the review, I will make the changes soon. Many of the reviews here also apply to the Flax Convnext PR so will make corresponding changes there too.

sanchit-gandhi

Awesome, thanks for iterating @Shubhamai! Really nice work on the clean-up! Almost ready for a third and final review 🤗 Then we can get it merged!

sanchit-gandhi · 2023-03-10T14:44:19Z

+    in_channels: int
+    out_channels: int
+    stride: int = 1
+    activation: Optional[str] = "relu"


Still unused!

sanchit-gandhi · 2023-03-10T14:47:20Z

+            padding=self.kernel_size // 2,
+            dtype=self.dtype,
+            use_bias=False,
+            kernel_init=nn.initializers.variance_scaling(2.0, mode="fan_out", distribution="normal", dtype=self.dtype),


Awesome - thanks for the explanation! Sounds good to me!

sanchit-gandhi · 2023-03-17T17:00:00Z

Hey @Shubhamai! Really nice work on this PR - just a few small changes to go now! Feel free to ping me once you're happy with the last bits and I'll get you a final review. Taking a look at Flax RegNet in the meantime!

Shubhamai · 2023-03-23T12:33:47Z

Made the request changes and apologies for the late reply, had a busy schedule lately.

sanchit-gandhi

Super! Thanks for iterating @Shubhamai! Requesting final review from @amyeroberts (maybe also @sgugger since it's the end of Amy's watch!)

amyeroberts

Thanks for your contribution and adding this model!

Left some tiny nits, but overall very nice and clean PR ❤️

amyeroberts · 2023-03-24T16:10:37Z

+RESNET_START_DOCSTRING = r"""
+
+    This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading, saving and converting weights from PyTorch models)


Suggested change

library implements for all its model (such as downloading, saving and converting weights from PyTorch models)

library implements for all its model (such as downloading, saving and converting weights from Flax models)

This should be PyTorch! We're converting from PyTorch models to Flax models - it's the same docstring we use in all Flax modelling files, e.g. for Flax BERT:

transformers/src/transformers/models/bert/modeling_flax_bert.py

Line 94 in 6587125

library implements for all its model (such as downloading, saving and converting weights from PyTorch models)

cc @Shubhamai

Ah, that makes sense, my bad 😅

No - my bad - sorry for the incorrect suggestion!

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

sanchit-gandhi · 2023-03-24T17:37:42Z

Thanks for the review @amyeroberts 🙌 Feel free to propagate the changes forward to Flax RegNet @Shubhamai and we can get a final review there!

sanchit-gandhi · 2023-03-24T17:40:13Z

We can merge when the CI is green! 🟢

Shubhamai · 2023-03-27T11:42:44Z

Thank you so much for the review & merge @sanchit-gandhi @amyeroberts , really appreciate your work & time on this ❤️

* [WIP] flax resnet * added pretrained flax models, results reproducible * Added pretrained flax models, results reproducible * working on tests * no real code change, just some comments * [flax] adding support for batch norm layers * fixing bugs related to pt+flax integration * removing loss from modeling flax output class * fixing classifier tests * fixing comments, model output * cleaning comments * review changes * review changes * Apply suggestions from code review Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * renaming Flax to PyTorch --------- Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

Shubhamai added 7 commits January 23, 2023 23:11

[WIP] flax resnet

4e6df72

added pretrained flax models, results reproducible

37116cd

Added pretrained flax models, results reproducible

7c2d042

working on tests

9387251

Merge branch 'huggingface:main' into resnet-flax

0f48e1f

no real code change, just some comments

f3def8f

Merge branch 'huggingface:main' into resnet-flax

46d869a

Shubhamai changed the title ~~Resnet flax~~ [WIP] Resnet flax Feb 6, 2023

[flax] adding support for batch norm layers

8799f17

Shubhamai mentioned this pull request Feb 11, 2023

[Flax] adding support for batch norm layers #21581

Merged

5 tasks

Shubhamai added 5 commits February 12, 2023 00:12

Merge branch 'main' into resnet-flax

bc996e8

fixing bugs related to pt+flax integration

b3ad809

removing loss from modeling flax output class

48b2217

fixing classifier tests

6dceb1e

Merge branch 'huggingface:main' into resnet-flax

88ac7a7

Shubhamai added 3 commits February 24, 2023 13:33

Merge branch 'main' into resnet-flax

e98f2aa

fixing comments, model output

a9cdbad

cleaning comments

e9467a8

Shubhamai changed the title ~~[WIP] Resnet flax~~ Resnet flax Feb 24, 2023

Shubhamai marked this pull request as ready for review February 24, 2023 09:23

sanchit-gandhi reviewed Feb 27, 2023

View reviewed changes

review changes

ed90f8b

Shubhamai mentioned this pull request Mar 9, 2023

Convnext flax #21485

Closed

8 tasks

sanchit-gandhi approved these changes Mar 10, 2023

View reviewed changes

review changes

7676b45

sanchit-gandhi approved these changes Mar 24, 2023

View reviewed changes

amyeroberts self-requested a review March 24, 2023 16:07

amyeroberts approved these changes Mar 24, 2023

View reviewed changes

Shubhamai and others added 2 commits March 24, 2023 22:35

Apply suggestions from code review

6dd6f2b

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

Merge branch 'huggingface:main' into resnet-flax

0f78c88

renaming Flax to PyTorch

c121172

amyeroberts merged commit a0cbbba into huggingface:main Mar 24, 2023

	library implements for all its model (such as downloading, saving and converting weights from PyTorch models)
	library implements for all its model (such as downloading, saving and converting weights from Flax models)

Uh oh!

Conversation

Shubhamai commented Feb 6, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Flax Implementation of microsoft/resnet-50

Before submitting

Who can review?

Status

TODO

Uh oh!

HuggingFaceDocBuilderDev commented Feb 6, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sanchit-gandhi commented Feb 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Shubhamai commented Feb 24, 2023

Uh oh!

sanchit-gandhi left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Shubhamai commented Feb 28, 2023

Uh oh!

sanchit-gandhi left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sanchit-gandhi commented Mar 17, 2023

Uh oh!

Shubhamai commented Mar 23, 2023

Uh oh!

sanchit-gandhi left a comment

Choose a reason for hiding this comment

Uh oh!

amyeroberts left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sanchit-gandhi Mar 24, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sanchit-gandhi commented Mar 24, 2023

Shubhamai commented Feb 6, 2023 •

edited

Loading

Flax Implementation of `microsoft/resnet-50`

HuggingFaceDocBuilderDev commented Feb 6, 2023 •

edited

Loading

sanchit-gandhi commented Feb 22, 2023 •

edited

Loading

sanchit-gandhi left a comment •

edited

Loading

sanchit-gandhi Mar 24, 2023 •

edited

Loading