Add DPT Flax #17779

younesbelkada · 2022-06-20T09:59:21Z

What does this PR do?

I tried to implement DPT (Dense Prediction with Transformers) in Flax during my free time! 🚀
By the way it is the first Segmentation and Depth Estimation model implemented in Flax on the library!

Nits/TODOs:

Figure out how to properly call BatchNorm and Dropout inside a Sequential
Deal correctly with Sequential layers
Test equivalency tests
Write documentation - For now they're just copy/pasted
Quetions:
Why the loss is not implemented in modeling_dpt.py ? I can probably help on that since I have already implemented the loss for a university project: https://github.com/antocad/FocusOnDepth/blob/master/FOD/Loss.py

cc @NielsRogge @sanchit-gandhi @patil-suraj

- added DPT in Flax - all non slow tests passes in local - still some nits have to be investigated

…e' object has no attribute 'tolist'`

- BN seems to work now - Equivalency test pass with tol=1e-4 but only with a hack

younesbelkada · 2022-06-20T10:00:53Z

All the keys match now but the equivalency test does not pass with 1e-5 but 1e-4 instead

ArthurZucker · 2022-06-20T10:01:26Z

src/transformers/modeling_flax_outputs.py

+
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):


Missing predicted depth documentation

src/transformers/models/dpt/modeling_flax_dpt.py

HuggingFaceDocBuilderDev · 2022-06-20T10:11:59Z

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

src/transformers/models/dpt/modeling_flax_dpt.py

ArthurZucker · 2022-06-20T14:59:37Z

src/transformers/models/dpt/modeling_flax_dpt.py

+class FlaxDPTUpsample(nn.Module):
+    scale: int = 2
+    method: str = "bilinear"
+
+    def setup(self):
+        pass
+
+    def __call__(self, x, output_size=None):
+        if output_size is None:
+            output_size = x.shape
+            output_size = (output_size[0], output_size[1] * self.scale, output_size[2] * self.scale, output_size[3])
+        return jax.image.resize(x, output_size, method="bilinear")


Should support align_corners = True

- more documentation - fix nit

NielsRogge · 2022-06-22T13:47:23Z

Would be great to also incorporate the updates of #17731

younesbelkada · 2022-06-22T14:01:20Z

src/transformers/models/dpt/modeling_flax_dpt.py

+
+    def setup(self):
+        self.cls_token = self.param("cls_token", nn.initializers.zeros, (1, 1, self.config.hidden_size))
+        self.patch_embeddings = FlaxPatchEmbeddings(self.config, dtype=self.dtype)


@NielsRogge I think that in this implementation we directly initialize the modules with the config (contrary than in DPTViTEmbeddings) if this is what you meant?

I copied these modules from FlaxViTModel which seem to have the right structure as suggested in #17731

Ok, then it's alright :)

Oh but the sanity check of the channel size is missing indeed, will add that!

Yeah we don't have it for other Flax models as well right now. Ideally (and also for consistency), we should have it for all vision models

- add custom conv transpose2d function - modify test

younesbelkada · 2022-06-25T18:13:59Z

The Flax model finally predicts the correct depths for the cats (left is Flax and right is Pytorch)!

For that it appears that the transpose conv does not give the same result as Pytorch's implementation that uses a gradient based operation. I fixed it by creating a custom function based on this PR: jax-ml/jax#5772 the PR does not seem to be merged soon. We can probably go for this hack for now until the PR in JAX gets merged

…to dpt-flax-younes

ArthurZucker · 2022-06-25T21:17:52Z

As wee discussed, it seems that align_corners set to False for both model would not require lowering the tolerance in one of the cases right?

younesbelkada · 2022-06-25T22:21:32Z

@ArthurZucker exact. I have put a new attribute in the DPTConfig and modified a bit the original modeling code but should not break backward compatibility. Now all tests pass with a tolerance of 1e-5

- added new attribute to config without breaking backward compatibility - modified a bit the tests

sanchit-gandhi

Hey @younesbelkada! Looks pretty good from the Flax side of things! Left a few requests, the overarching one being the use of # Copied from... statements, both for internal Transformers code and that copied externally (e.g. from Haiku). Really helps in knowing what are the salient portions of the modelling code to review! But otherwise a very strong effort on the Flax front 💪

src/transformers/models/dpt/configuration_dpt.py

src/transformers/models/dpt/gradient_convolution.py

src/transformers/models/dpt/modeling_flax_dpt.py

github-actions · 2022-07-26T15:02:06Z

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.

Co-authored-by: Sanchit Gandhi <[email protected]>

- removed `copied_from` on non module objects - check why `FlaxDPTViTLayerCollection` is not copied from `FlaxViTLayerCollection`

- added correct link for `CopiedFrom` - Added explicit argument for transposed conv on model def

younesbelkada · 2022-08-10T18:20:13Z

Thank you very much @sanchit-gandhi for the very detailed review! I had a second round of refactoring while catching up on Flax projects and would love to have a second round of review (left also some unresolved comments) 💪 Thanks again 🙏

github-actions · 2022-10-07T15:04:08Z

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.

younesbelkada and others added 5 commits June 18, 2022 22:34

first commit:

1166511

- added DPT in Flax - all non slow tests passes in local - still some nits have to be investigated

make quality

563217b

implement pt_flax equivlence test to bypass `AttributeError: 'NoneTyp…

77220f7

…e' object has no attribute 'tolist'`

add few changes:

5f1341e

- BN seems to work now - Equivalency test pass with tol=1e-4 but only with a hack

add modification sequential

86320ad

younesbelkada requested a review from ArthurZucker June 20, 2022 09:59

ArthurZucker reviewed Jun 20, 2022

View reviewed changes

src/transformers/models/dpt/modeling_flax_dpt.py Outdated Show resolved Hide resolved

younesbelkada marked this pull request as ready for review June 20, 2022 10:54

fix copies

6b2c28c

younesbelkada changed the title ~~Add Flax implementation of DPT~~ [WIP] Add Flax implementation of DPT Jun 20, 2022

ArthurZucker reviewed Jun 20, 2022

View reviewed changes

src/transformers/models/dpt/modeling_flax_dpt.py Show resolved Hide resolved

ArthurZucker reviewed Jun 20, 2022

View reviewed changes

younesbelkada mentioned this pull request Jun 20, 2022

Adding align_corners to jax.image.resize google/flax#2210

Closed

few fixes

926fb77

- more documentation - fix nit

younesbelkada mentioned this pull request Jun 21, 2022

Adding align_corners to jax.image.resize google/flax#2214

Closed

younesbelkada commented Jun 22, 2022

View reviewed changes

younesbelkada mentioned this pull request Jun 22, 2022

Adding align_corners to jax.image.resize jax-ml/jax#11206

Open

changes

72bf6a4

- add custom conv transpose2d function - modify test

ArthurZucker force-pushed the dpt-flax-younes branch from ccfb913 to a379a2f Compare June 25, 2022 20:23

Merge branch 'main' of https://github.com/huggingface/transformers in…

7762a1e

…to dpt-flax-younes

ArthurZucker force-pushed the dpt-flax-younes branch from a379a2f to 7762a1e Compare June 25, 2022 20:25

ArthurZucker added 4 commits June 25, 2022 22:42

style

63a6144

add FlaxViTPatchEmbeddings for consistency

3d74f94

update tests

ed7a4bc

consistency

e2a61c9

style

e11b043

all tests should pas

05dcc85

- added new attribute to config without breaking backward compatibility - modified a bit the tests

younesbelkada changed the title ~~[WIP] Add Flax implementation of DPT~~ Add Flax implementation of DPT Jun 25, 2022

younesbelkada requested review from NielsRogge and sanchit-gandhi June 25, 2022 22:49

sanchit-gandhi suggested changes Jul 1, 2022

View reviewed changes

github-actions bot closed this Aug 3, 2022

younesbelkada reopened this Aug 3, 2022

younesbelkada and others added 7 commits August 10, 2022 15:26

fixing few comments

7c26f69

Apply suggestions from code review

3505582

Co-authored-by: Sanchit Gandhi <[email protected]>

Update src/transformers/models/dpt/gradient_convolution.py

425508e

Co-authored-by: Sanchit Gandhi <[email protected]>

add few comments

24aeb4d

Apply suggestions from code review

479d0e8

Co-authored-by: Sanchit Gandhi <[email protected]>

refactor a bit:

791ea24

- removed `copied_from` on non module objects - check why `FlaxDPTViTLayerCollection` is not copied from `FlaxViTLayerCollection`

Merge remote-tracking branch 'upstream/main' into dpt-flax-younes

010a915

younesbelkada changed the title ~~Add Flax implementation of DPT~~ Add DPT Flax Aug 10, 2022

younesbelkada added 2 commits August 10, 2022 18:49

add comments on key naming strategy

a40f21d

few modifications

f945eef

- added correct link for `CopiedFrom` - Added explicit argument for transposed conv on model def

younesbelkada requested a review from sanchit-gandhi August 10, 2022 18:20

github-actions bot closed this Sep 12, 2022

huggingface deleted a comment from github-actions bot Sep 12, 2022

sanchit-gandhi reopened this Sep 12, 2022

younesbelkada mentioned this pull request Sep 17, 2022

Add from_pt argument in .from_pretrained huggingface/diffusers#527

Merged

github-actions bot closed this Oct 16, 2022


		hidden_states (`tuple(torch.FloatTensor)`, optional, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):

Add DPT Flax #17779

Add DPT Flax #17779

Uh oh!

Conversation

younesbelkada commented Jun 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Uh oh!

younesbelkada commented Jun 20, 2022

Uh oh!

ArthurZucker Jun 20, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Jun 20, 2022

Uh oh!

Uh oh!

ArthurZucker Jun 20, 2022

Choose a reason for hiding this comment

Uh oh!

NielsRogge commented Jun 22, 2022

Uh oh!

younesbelkada Jun 22, 2022 • 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!

younesbelkada Jun 22, 2022

Choose a reason for hiding this comment

Uh oh!

NielsRogge Jun 22, 2022

Choose a reason for hiding this comment

Uh oh!

younesbelkada Jun 22, 2022

Choose a reason for hiding this comment

Uh oh!

NielsRogge Jun 22, 2022

Choose a reason for hiding this comment

Uh oh!

younesbelkada commented Jun 25, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ArthurZucker commented Jun 25, 2022

Uh oh!

younesbelkada commented Jun 25, 2022

Uh oh!

sanchit-gandhi left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Jul 26, 2022

Uh oh!

younesbelkada commented Aug 10, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Oct 7, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

younesbelkada commented Jun 20, 2022 •

edited

Loading

younesbelkada Jun 22, 2022 •

edited

Loading

younesbelkada commented Jun 25, 2022 •

edited

Loading

younesbelkada commented Aug 10, 2022 •

edited

Loading