Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

TensorFlow MobileViT #18555

Merged
merged 26 commits into from
Sep 1, 2022
Merged

TensorFlow MobileViT #18555

merged 26 commits into from
Sep 1, 2022

Conversation

sayakpaul
Copy link
Member

@sayakpaul sayakpaul commented Aug 10, 2022

This PR implements the MobileViT model in TensorFlow.

Interesting points

  • The classification and segmentation models provided with MobileViT are fully compatible with TensorFlow Lite. Therefore, I have included sample code in the model documentation showing how to perform the TensorFlow Lite conversion (~4 lines of code).
  • TFLlite versions of the smallest checkpoints for classification and semantic segmentation are 1MB and 2MBs, respectively. I believe this will be quite beneficial to the TinyML community.

TODOs

  • Hosting of the TF checkpoints on the Hub. (Can I do it now? If so, I need resources that show how to do that.)
  • Remove from_pt wherever needed.

@amyeroberts @gante @sgugger up for review!

@HuggingFaceDocBuilderDev
Copy link

HuggingFaceDocBuilderDev commented Aug 10, 2022

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

@sayakpaul sayakpaul changed the title [WIP] TensorFlow MobileViT TensorFlow MobileViT Aug 22, 2022
@sayakpaul sayakpaul marked this pull request as ready for review August 22, 2022 07:10
@gante gante requested review from sgugger, amyeroberts and gante August 24, 2022 14:56
Copy link
Member

@gante gante left a comment

Choose a reason for hiding this comment

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

Thank you for the TF implementation of MobileViT 🙌 The TFLite demo is great, especially because it should be covered by our doctests 🚀

In addition to the comments throughout the code, I have the following notes:

  1. I will share the instructions to open the Hub PR when this PR is approved by all (everyone has permission to do it now 🎉 )
  2. The training argument is missing in the layer's call (and in places like the dropout calls)

@sayakpaul sayakpaul requested review from gante and removed request for amyeroberts and sgugger August 25, 2022 04:45
@sayakpaul
Copy link
Member Author

Didn't realize that re-requesting a review from @gante would result in removing @amyeroberts and @sgugger from the reviewer list. Please know that it was completely unintentional.

@gante
Copy link
Member

gante commented Aug 25, 2022

@sayakpaul no worries :)

@gante gante requested review from sgugger and amyeroberts August 25, 2022 10:23
Copy link
Member

@gante gante left a comment

Choose a reason for hiding this comment

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

Thank you for the changes 🔥

Copy link
Collaborator

@amyeroberts amyeroberts left a comment

Choose a reason for hiding this comment

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

LGTM! 📱

src/transformers/models/mobilevit/__init__.py Outdated Show resolved Hide resolved
src/transformers/models/mobilevit/modeling_tf_mobilevit.py Outdated Show resolved Hide resolved
patch_width, patch_height = self.patch_width, self.patch_height
patch_area = tf.cast(patch_width * patch_height, "int32")

batch_size, orig_height, orig_width, channels = shape_list(features)
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
batch_size, orig_height, orig_width, channels = shape_list(features)
batch_size, orig_height, orig_width, channels = tf.shape(features)

Copy link
Member Author

Choose a reason for hiding this comment

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

Having it in one line leads to:

OperatorNotAllowedInGraphError: Iterating over a symbolic tf.Tensor is not allowed in Graph execution. Use Eager execution or decorate this function with @tf.function.

That's why I separated it.

def test_attention_outputs(self):
pass

@unittest.skip("Test was written for TF 1.x and isn't really relevant here")
Copy link
Collaborator

Choose a reason for hiding this comment

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

If this is the case - should it even be in the test suite? cc @gante

Copy link
Member

Choose a reason for hiding this comment

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

Nope, me and @Rocketknight1 talked about it a few weeks ago. We should remove this test, it's heavy and the only new thing it tests is that we can build a functional TF model with the model class (which it's kinda obvious we can)

Copy link
Member Author

@sayakpaul sayakpaul Aug 26, 2022

Choose a reason for hiding this comment

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

I assume it will be phased out in a separate PR from the main TF testing suite?

Copy link
Member

Choose a reason for hiding this comment

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

Ofc, as a separate PR 👍 Leave it be as it is in this PR :)

src/transformers/models/mobilevit/modeling_tf_mobilevit.py Outdated Show resolved Hide resolved
src/transformers/models/mobilevit/modeling_tf_mobilevit.py Outdated Show resolved Hide resolved
@amyeroberts
Copy link
Collaborator

Thanks for another great model addition @sayakpaul !

@gante
Copy link
Member

gante commented Aug 25, 2022

@sayakpaul assuming it is passing the slow tests, it is ready for the TF weights.

The super complex instructions to do it are as follows:

  1. Make sure you have the latest version of the hub installed (pip install huggingface_hub -U) and that you are logged in to HF with a write token (huggingface-cli login)
  2. Run transformers-cli pt-to-tf --model-name foo/bar from this branch :D
  3. In the Hub PR, tag @joaogante, @lysandre

@sayakpaul
Copy link
Member Author

Super simple (complex?) question:

What is the format of foo/bar?

@gante
Copy link
Member

gante commented Aug 25, 2022

The same as the model name on the hub, e.g. this model would be apple/mobilevit-small

P.S.: I edited the comment above with a 3rd step :D

@hollance
Copy link
Contributor

[...] the output consistency of nn.functional.interpolate and tf.image.resize with the same argument values.

This might be due to the align_corners option. I once wrote a long blog post about this difference between PyTorch and TF. https://machinethink.net/blog/coreml-upsampling/ Not sure if that's the same issue but it seems likely.

@sayakpaul
Copy link
Member Author

[...] the output consistency of nn.functional.interpolate and tf.image.resize with the same argument values.

This might be due to the align_corners option. I once wrote a long blog post about this difference between PyTorch and TF. https://machinethink.net/blog/coreml-upsampling/ Not sure if that's the same issue but it seems likely.

Very well! If we need to deal with the inconsistencies between tf.image.resize and nn.functional.interpolate I suggest we do that in a separate PR 'cause various vision models would benefit from that (ViT for example).

@sayakpaul
Copy link
Member Author

@gante WDYT?

@gante
Copy link
Member

gante commented Aug 30, 2022

@sayakpaul regarding the PR, all good on my end, but we still need approval from @sgugger :D

As for the tf.image.resize -- yeah, it would be nice to standardize for all models. Would you be interested in working on it? In any case, I'd like to ask you to open an issue, so we don't forget to track it!

@sayakpaul
Copy link
Member Author

As for the tf.image.resize -- yeah, it would be nice to standardize for all models. Would you be interested in working on it? In any case, I'd like to ask you to open an issue, so we don't forget to track it!

On it, sir!

@sayakpaul
Copy link
Member Author

@amyeroberts @gante

Please take note of the changes in 32cfd30.

Initially, when I tested TFLite conversion it didn't require any spec for SELECT operations but now they're failing with a specification for the SELECT ops. What is more surprising is that the TFLite interpreter is treating tf.Conv2D to be a SELECT op. Hence I have raised tensorflow/tensorflow#57550.

Copy link
Collaborator

@sgugger sgugger left a comment

Choose a reason for hiding this comment

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

Thanks a lot for your PR! Left a couple of nits then we can merge this.

sayakpaul and others added 6 commits September 1, 2022 16:59
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
@gante
Copy link
Member

gante commented Sep 1, 2022

(retriggered failing job, seems like a spurious failure)

@sayakpaul
Copy link
Member Author

Yeah probably nothing related to the PR?

@sgugger
Copy link
Collaborator

sgugger commented Sep 1, 2022

The build doc job failure is not spurious. There seems to be a problem with an example bloc introduced by this PR.

@sayakpaul
Copy link
Member Author

Let me see if removing comments from the example block does the trick. Because when the job wasn't failing the example block didn't have any comments.

@sayakpaul
Copy link
Member Author

No, it didn't help :( Any suggestions to try out?

You can use the following code to convert a MobileViT checkpoint (be it image classification or semantic segmentation) to generate a
TensorFlow Lite model:

```py
Copy link
Member

Choose a reason for hiding this comment

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

could it be because the example is indented?

Copy link
Collaborator

Choose a reason for hiding this comment

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

Looks like it, for some reason. The failure seems to disappear locally when I remove it. In any case its place is probably closer to the TF models doc?

Copy link
Member Author

Choose a reason for hiding this comment

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

This was the culprit it seems :3

@gante
Copy link
Member

gante commented Sep 1, 2022

The build doc job failure is not spurious. There seems to be a problem with an example bloc introduced by this PR.

My bad :D read the failure bottom to top, so I didn't notice the mobilevit errors

@sgugger sgugger merged commit 954e18a into huggingface:main Sep 1, 2022
@sayakpaul sayakpaul deleted the feat/tf-mobilevit branch September 1, 2022 15:12
oneraghavan pushed a commit to oneraghavan/transformers that referenced this pull request Sep 26, 2022
* initial implementation.

* add: working model till image classification.

* add: initial implementation that passes intg tests.

Co-authored-by: Amy <aeroberts4444@gmail.com>

* chore: formatting.

* add: tests (still breaking because of config mismatch).

Coo-authored-by: Yih <2521628+ydshieh@users.noreply.github.com>

* add: corrected tests and remaning changes.

* fix code style and repo consistency.

* address PR comments.

* address Amy's comments.

* chore: remove from_pt argument.

* chore: add full-stop.

* fix: TFLite model conversion in the doc.

* Update src/transformers/models/mobilevit/modeling_tf_mobilevit.py

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

* Update src/transformers/models/mobilevit/modeling_tf_mobilevit.py

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

* Update src/transformers/models/mobilevit/modeling_tf_mobilevit.py

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

* Update src/transformers/models/mobilevit/modeling_tf_mobilevit.py

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

* Update src/transformers/models/mobilevit/modeling_tf_mobilevit.py

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

* apply formatting.

* chore: remove comments from the example block.

* remove identation in the example.

Co-authored-by: Amy <aeroberts4444@gmail.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
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.

6 participants