Convnext flax by Shubhamai · Pull Request #21485 · huggingface/transformers

Shubhamai · 2023-02-07T08:24:31Z

Flax Implementation of `facebook/convnext-tiny-224`

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?

Flax: @sanchit-gandhi

TODO

Last Updated : 10 Feb, 2023

Fixing tests failed in ci/circleci: tests_flax actions.
Uploading Shubhamai/convnext-tiny-224 flax weights to facebook/convnext-tiny-224.
Depends on merge of Improve flax test_attention_outputs PR to pass (or technically skip) test_attention_outputs test.

HuggingFaceDocBuilderDev · 2023-02-07T08:39:46Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

Shubhamai · 2023-02-23T17:53:33Z

@sanchit-gandhi this PR is also ready for your review in case it was missed. And again, thank you so much for taking the time to review the PR.

Shubhamai · 2023-03-09T11:18:56Z

@sanchit-gandhi Reminder incase my previous message got missed! Also the #21472 ( previous reviews implemented ) and #21867 PR are awaiting review. Thanks!

sanchit-gandhi

Looks great already @Shubhamai! Think we just need to propagate some of the changes from Flax ResNet / RegNet into this PR, then we're good to go 🚀

sanchit-gandhi · 2023-04-18T16:17:56Z

            The initial value for the layer scale.
        drop_path_rate (`float`, *optional*, defaults to 0.0):
            The drop rate for stochastic depth.
+        image_size (`int`, *optional*, defaults to 224):


Think we can just remove this from the docstring (see #21472 (comment))

sanchit-gandhi · 2023-04-18T16:19:42Z

+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(self, hidden_states: jnp.ndarray, train: bool = False) -> jnp.ndarray:


We can also swap train for deterministic (see #21472 (comment))

Ah, missed it there, should be fixed in recent commit

sanchit-gandhi · 2023-04-18T16:21:41Z

+        shape = (hidden_states.shape[0],) + (1,) * (
+            hidden_states.ndim - 1
+        )  # work with diff dim tensors, not just 2D ConvNets
+        random_tensor = keep_prob + random.uniform(random.PRNGKey(0), shape, dtype=self.dtype)


We probably need to pass a PRNG key here (rather than always starting from seed 0), see

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

Lines 364 to 366 in 90247d3

dropout_rng = None

if not deterministic and self.config.attention_probs_dropout_prob > 0.0:

dropout_rng = self.make_rng("dropout")

and #21023 (review)

Replaced random.PRNGKey(0) with self.make_rng("dropout")

sanchit-gandhi · 2023-04-18T16:23:23Z

+
+    def __call__(self, pixel_values: jnp.ndarray) -> jnp.ndarray:
+        num_channels = pixel_values.shape[-1]
+        if num_channels != self.config.num_channels:


Same here regarding the config and num_channels

I assume you mean removing the if statement and num_channels in the config file. If yes, then the changed should be made in the recent commit

Sorry we shouldn't remove num_channels from the config as this will break the PyTorch code!

sanchit-gandhi · 2023-04-18T16:32:43Z

+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if pixel_values is None:


This won't play very nicely with JIT since we have control flow on our input (see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#control-flow)

Think we can safely assume here that pixel_values will not be None (otherwise the next line self.embeddings(pixel_values) will complain anyway that pixel_values is None), and so can remove this check.

Agree, should be removed in the recent commit

sanchit-gandhi · 2023-04-18T16:33:32Z

+    >>> image = Image.open(requests.get(url, stream=True).raw)
+
+    >>> image_processor = AutoImageProcessor.from_pretrained("facebook/convnext-tiny-224")
+    >>> model = FlaxConvNextModel.from_pretrained("facebook/convnext-tiny-224")


Cool! Have we uploaded the Flax weights already?

Just made the PR https://huggingface.co/facebook/convnext-tiny-224/discussions/3

sanchit-gandhi · 2023-04-18T16:34:01Z

+    from transformers import AutoFeatureExtractor
+
+
+class FlaxConvNextModelTester(unittest.TestCase):


Would be awesome to update the tester according to the Flax ResNet model tester :)

sanchit-gandhi · 2023-04-18T16:34:25Z

+
+    @slow
+    def test_inference_image_classification_head(self):
+        model = FlaxConvNextForImageClassification.from_pretrained("Shubhamai/convnext-tiny-224")


Feel free to open a PR for the Flax weights on the Facebook org!

PR https://huggingface.co/facebook/convnext-tiny-224/discussions/3

sanchit-gandhi · 2023-05-05T15:44:10Z


    def __init__(
        self,
-        num_channels=3,


We can't remove config.num_channels since this will break the PyTorch modelling code! https://github.com/huggingface/transformers/blob/40082d598b756c4f1fb048c571ed95c31f05a69c/src/transformers/models/convnext/modeling_convnext.py#LL137C2-L137C2

Hi, I should have mentioned the review changes weren't fully complete, I am aware of the error due to failed tests, I will fix it asap ( probably on this weekend )

Cool! Sounds good! Feel free to ping me when it's ready and I'll get you a final review :)

sanchit-gandhi · 2023-05-05T15:46:59Z

+
+    def __call__(self, pixel_values: jnp.ndarray) -> jnp.ndarray:
+        num_channels = pixel_values.shape[-1]
+        if num_channels != self.config.num_channels:


Sorry we shouldn't remove num_channels from the config as this will break the PyTorch code!

sanchit-gandhi · 2023-05-05T15:48:10Z

        layer_norm_eps=1e-12,
        layer_scale_init_value=1e-6,
        drop_path_rate=0.0,
-        image_size=224,


In general there shouldn't be a need to change the config file at all for a Flax port

github-actions · 2023-05-30T15:02:51Z

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

Shubhamai added 6 commits February 2, 2023 20:16

flax implementation of convnext, some tests left

a9372ab

Merge branch 'huggingface:main' into convnext-flax

8469640

fixed some tests, 2 still failing

ee7848a

formatting code

ae42181

fixing tests

747762f

Merge branch 'huggingface:main' into convnext-flax

3296fce

Shubhamai mentioned this pull request Feb 7, 2023

[Tests] Improve flax test_attention_outputs #21486

Merged

5 tasks

Shubhamai added 4 commits February 10, 2023 15:33

Merge branch 'huggingface:main' into convnext-flax

1a64729

fixing pytorch-flax tests

aef4068

Merge branch 'huggingface:main' into convnext-flax

4833573

removing loss from modeling flax output class

f3da9b3

Shubhamai changed the title ~~[WIP] Convnext flax~~ Convnext flax Feb 10, 2023

Shubhamai marked this pull request as ready for review February 11, 2023 18:12

Shubhamai mentioned this pull request Feb 14, 2023

[Flax] adding support for batch norm layers #21581

Merged

5 tasks

Shubhamai mentioned this pull request Feb 28, 2023

Resnet flax #21472

Merged

7 tasks

Shubhamai added 2 commits February 28, 2023 21:18

[resnet pr] review changes

1880172

Merge branch 'huggingface:main' into convnext-flax

0f354a3

huggingface deleted a comment from github-actions Bot Apr 4, 2023

Merge branch 'main' into convnext-flax

5eb91be

sanchit-gandhi reviewed Apr 18, 2023

View reviewed changes

Shubhamai added 2 commits April 21, 2023 20:58

Merge branch 'huggingface:main' into convnext-flax

cd3820e

review changes

14a1e0e

sanchit-gandhi reviewed May 5, 2023

View reviewed changes

github-actions Bot closed this Jun 7, 2023

	dropout_rng = None
	if not deterministic and self.config.attention_probs_dropout_prob > 0.0:
	dropout_rng = self.make_rng("dropout")

		from transformers import AutoFeatureExtractor


		class FlaxConvNextModelTester(unittest.TestCase):

Uh oh!

Conversation

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

Uh oh!

Flax Implementation of facebook/convnext-tiny-224

Before submitting

Who can review?

TODO

Uh oh!

HuggingFaceDocBuilderDev commented Feb 7, 2023

Uh oh!

Shubhamai commented Feb 23, 2023

Uh oh!

Shubhamai commented Mar 9, 2023

Uh oh!

sanchit-gandhi left a comment

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!

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!

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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Shubhamai May 5, 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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented May 30, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Shubhamai commented Feb 7, 2023 •

edited

Loading

Flax Implementation of `facebook/convnext-tiny-224`

Shubhamai May 5, 2023 •

edited

Loading