[DeepVision Port] SegFormer and Mix-Transformers #1946

DavidLandup0 · 2023-07-13T18:15:56Z

What does this PR do?

As discussed in #1933 - setting up a draft PR for porting SegFormer and associated layers into KCV. Draft PR for now with placeholder main model dump, layers and tests incoming soon. Will tag once ready for review.

Demo Notebooks

Mix-Transformer components, from_preset() usage and training: https://colab.research.google.com/drive/1Q3m9-LKICrFzuUhVMIPd7pY2l9Z3BLhg?usp=sharing
SegFormer head, from_preset() usage and training: #TODO

Questions and API Considerations

DeepLabV3 takes any backbone, but SegFormer is meant to be used with MiT (Mix Transformers), and depends on the output channels which is a field defined in the model. Should we make it generally usable with other backbones? IMO, no, since the head is really just an MLP head, and the crux of the paper is MiT.
There's no name for the type of attention they use, but they refer to it as efficient attention in the paper. What name should we use? SegFormerMultiHeadAttention sounds like a mouthful.
How do we expose the API if we don't support a backbone argument? Just SegFormer.from_preset()?

Before submitting

Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue? Please add a link
to it if that's the case. Porting DeepVision into KerasCV #1933
Did you write any new necessary tests?
If this adds a new model, can you run a few training steps on TPU in Colab to ensure that no XLA incompatible OP are used?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ianstenbit just tagging so you can follow the progress as it comes in. Otherwise, no need to spend time until it's un-drafted for review :)

DavidLandup0 · 2023-07-17T12:35:16Z

Just to update you @ianstenbit - the port is going well, but it took me a bit longer than anticipated to get used to Keras Core + the new API 😅

I ran into a small blocker and documented it here since I'm not sure what the intended usage is when returning tensors and non-tensors from a call(). If you've encountered this before, any idea for a workaround would be greatly appreciated 🙇

ianstenbit

Thanks David! I'm taking a look at the non-tensor return issue.

We have a few options for workarounds, including:

Computing the shape outside of the layer (because from my cursory view the returned shape is just the input shape / stride)
Making the PatchingEmbeddingLayer offer a new method which computes these values which callers can use.

That said, I think we should be able to make this work. I'm taking a look at your issue on Keras Core to see if I can get a working fix.

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py

DavidLandup0 · 2023-07-17T15:50:03Z

Thanks! I opened it as an issue since I'm not sure if it's the intended usage. If so, I'd go with computing the values outside the layer/with an extra method

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone_presets.py

keras_cv/layers/overlapping_patching_embedding.py

ianstenbit · 2023-07-17T16:12:14Z

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py

+class MiTBackbone(Backbone):
+    def __init__(
+        self,
+        input_shape=None,


Default to (None, None, 3) so that channel dims can be known at build time for conv layers

This'll have to default to (224, 224, 3) actually, since the input shape will have to be known at instantiation time

DavidLandup0 · 2023-07-23T21:24:29Z

@ianstenbit looks like MiTs are shaped up. Here's a demo notebook showcasing the components, inputs/output shapes, pyramid levels, from_preset() usage and training MiTs on a classification task: https://colab.research.google.com/drive/1Q3m9-LKICrFzuUhVMIPd7pY2l9Z3BLhg?usp=sharing

There are a couple of weird-looking ops.cast() calls that aren't very clean, and a custom reshaping layer since keras.Reshape() caused errors for some reason. I'd like to clean these up and sync up on whether there's a cleaner alternative for them :)

99% of the work are MiTs - SegFormer is just MiT+seg top. Could you please review the backbone while I shape up SegFormers? With a green light, I'll write up the unit tests and add proper docstrings.

ianstenbit

Generally looks good! Left you a few minor comments.

keras_cv/layers/efficient_multihead_attention.py

keras_cv/layers/hierarchical_transformer_encoder.py

keras_cv/layers/overlapping_patching_embedding.py

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py

keras_cv/models/segmentation/segformer/segformer.py

ianstenbit

Just a few things to clean up -- in the meantime I am seeing if the tests need any fixing.

keras_cv/layers/hierarchical_transformer_encoder.py

ianstenbit · 2023-08-21T18:18:02Z

keras_cv/models/segmentation/segformer/segformer_aliases.py

+class SegFormerB0(SegFormer):
+    def __new__(
+        cls,
+        num_classes=19,


We shouldn't specify a default for num_classes as it needs to be user-specified. A silent default could be very confusing.

Removed. Should it be requested as a mandatory arg?

Yes I think this should be required at init time.

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py

ianstenbit · 2023-08-21T20:38:56Z

/gcbrun

ianstenbit · 2023-08-21T20:53:17Z

Looks like tests are passing locally, but on CI (for TF) it will depend on us getting a new release of Keras Core which includes keras-team/keras-core#722

In the meantime, @DavidLandup0 I left a few review comments for you to take a look at -- thanks!

DavidLandup0 · 2023-08-22T20:10:20Z

Awesome, thanks! Getting to these soon. Thanks for the review pass! :)

keras_cv/layers/overlapping_patching_embedding.py

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py

ianstenbit

Thank you David!

I think this PR is basically all set, I just need to merge #2037 to fix CI

ianstenbit · 2023-08-24T20:35:25Z

/gcbrun

DavidLandup0 · 2023-08-24T21:20:45Z

Awesome, thank you for the help in the final stretch! @ianstenbit 🎉

ianstenbit · 2023-08-24T22:22:15Z

It looks like the GCB failure is because I need to update the Docker image of our GCB runners to use the newest Keras Core version -- doing that now.

ianstenbit · 2023-08-24T22:56:36Z

CI failures are unrelated -- seems like DeepLab + YOLOV8 have some breakages with the latest Keras Core version. I'll open a separate PR for those.

DavidLandup0 · 2023-08-24T23:04:31Z

Need a hand with DLV3 or YOLO?

ianstenbit · 2023-08-24T23:09:14Z

You're welcome to look if you'd like -- for DeepLab it's a deserialization issue. Haven't looked at YOLO yet.

You can repro by installing latest Keras Core version and running the large tests of those models with TF backend.
edit: probably best to just work on CLIP instead -- I can handle this part, it's not very fun anyway!

DavidLandup0 · 2023-08-24T23:32:50Z

Sure! Sign me up for YOLO if it's not too urgent then :)

* initial dump * add all basic layers, port roughly to keras core ops * updated .gitignore * segformer head and formatting * cleanup * remove tf call * remove tf * migrating to more keras ops * cleanups and fixes * fix reshaping * comments * from presets api, keras.ops -> ops * embed_dims -> embedding_dims * addressing some PR comments * docstrings, argument update * depths arg * sync * compute output shapes * segformer progress * head * softmax * remove softmax * undo compute_output_shapes() * efficientmultiheadattention -> segformermultiheadattention * docstrings * softmax output * segformer presets * updating segformer presets * segformer presets * import aliases * refactoring * pr comments * pr comments * add aliases * aliases ot init * refactor fix * import keras_cv_export * fix presets/aliases and add copyright * linter warnings * linter errors * consistency in presets * return config * fix serialization * Some cleanup + more tests * Fix DropPath layer (need to update tests + add shim for tf.keras * Finish DropPath layer * Use static shape in backbone * Formatting * Switch back to ops.shape * documentation * documentation * remove default num classes * fix docs --------- Co-authored-by: ianjjohnson <3072903+ianstenbit@users.noreply.github.com>

DavidLandup0 added 8 commits July 13, 2023 20:12

initial dump

dc41892

add all basic layers, port roughly to keras core ops

e5677e6

updated .gitignore

7bd1056

segformer head and formatting

03470df

cleanup

cb1c702

remove tf call

22f8fdf

remove tf

5c9803a

migrating to more keras ops

314dc6b

ianstenbit reviewed Jul 17, 2023

View reviewed changes

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py Outdated Show resolved Hide resolved

ianstenbit reviewed Jul 17, 2023

View reviewed changes

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py Show resolved Hide resolved

ianstenbit reviewed Jul 17, 2023

View reviewed changes

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone_presets.py Show resolved Hide resolved

ianstenbit reviewed Jul 17, 2023

View reviewed changes

keras_cv/layers/overlapping_patching_embedding.py Outdated Show resolved Hide resolved

ianstenbit reviewed Jul 17, 2023

View reviewed changes

DavidLandup0 added 5 commits July 23, 2023 21:58

cleanups and fixes

7a0151b

fix reshaping

44f01af

comments

eb5b5ae

from presets api, keras.ops -> ops

ea0239f

embed_dims -> embedding_dims

b6128a5

ianstenbit suggested changes Jul 24, 2023

View reviewed changes

DavidLandup0 added 8 commits July 24, 2023 16:56

addressing some PR comments

8322109

docstrings, argument update

75bb4a2

depths arg

97daf7c

sync

5f9dc0c

compute output shapes

efbbd49

segformer progress

d3b43c6

head

dab4e74

softmax

1dba059

ianstenbit suggested changes Aug 21, 2023

View reviewed changes

ianstenbit reviewed Aug 21, 2023

View reviewed changes

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py Show resolved Hide resolved

ianstenbit added 4 commits August 21, 2023 13:11

Some cleanup + more tests

eea5e3c

Fix DropPath layer (need to update tests + add shim for tf.keras

8e62cf6

Finish DropPath layer

b9efeb1

Use static shape in backbone

bd5a99f

ianstenbit added 2 commits August 21, 2023 14:41

Formatting

3d29b0a

Switch back to ops.shape

4e2c4e8

DavidLandup0 added 3 commits August 23, 2023 21:05

documentation

b32e0cf

documentation

743a3bb

remove default num classes

c640fc9

ianstenbit reviewed Aug 23, 2023

View reviewed changes

keras_cv/layers/overlapping_patching_embedding.py Outdated Show resolved Hide resolved

ianstenbit reviewed Aug 23, 2023

View reviewed changes

keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py Outdated Show resolved Hide resolved

fix docs

f1b5ffa

ianstenbit approved these changes Aug 23, 2023

View reviewed changes

Merge branch 'master' into segformer_tf

e32704b

ianstenbit merged commit ab812d1 into keras-team:master Aug 24, 2023
8 of 9 checks passed

sachinprasadhs mentioned this pull request Feb 15, 2024

Porting DeepVision into KerasCV #1933

Closed

DavidLandup0 mentioned this pull request Sep 29, 2024

[Mix Transformer] Add Presets for MiTB0...MiTB5 keras-team/keras-hub#1893

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[DeepVision Port] SegFormer and Mix-Transformers #1946

[DeepVision Port] SegFormer and Mix-Transformers #1946

DavidLandup0 commented Jul 13, 2023 •

edited

Loading

DavidLandup0 commented Jul 17, 2023 •

edited

Loading

ianstenbit left a comment

DavidLandup0 commented Jul 17, 2023

ianstenbit Jul 17, 2023

DavidLandup0 Jul 23, 2023

DavidLandup0 commented Jul 23, 2023 •

edited

Loading

ianstenbit left a comment

ianstenbit left a comment

ianstenbit Aug 21, 2023

DavidLandup0 Aug 23, 2023

ianstenbit Aug 23, 2023

ianstenbit commented Aug 21, 2023

ianstenbit commented Aug 21, 2023

DavidLandup0 commented Aug 22, 2023

ianstenbit left a comment

ianstenbit commented Aug 24, 2023

DavidLandup0 commented Aug 24, 2023

ianstenbit commented Aug 24, 2023

ianstenbit commented Aug 24, 2023

DavidLandup0 commented Aug 24, 2023

ianstenbit commented Aug 24, 2023 •

edited

Loading

DavidLandup0 commented Aug 24, 2023

[DeepVision Port] SegFormer and Mix-Transformers #1946

[DeepVision Port] SegFormer and Mix-Transformers #1946

Conversation

DavidLandup0 commented Jul 13, 2023 • edited Loading

What does this PR do?

Demo Notebooks

Questions and API Considerations

Before submitting

Who can review?

DavidLandup0 commented Jul 17, 2023 • edited Loading

ianstenbit left a comment

Choose a reason for hiding this comment

DavidLandup0 commented Jul 17, 2023

ianstenbit Jul 17, 2023

Choose a reason for hiding this comment

DavidLandup0 Jul 23, 2023

Choose a reason for hiding this comment

DavidLandup0 commented Jul 23, 2023 • edited Loading

ianstenbit left a comment

Choose a reason for hiding this comment

ianstenbit left a comment

Choose a reason for hiding this comment

ianstenbit Aug 21, 2023

Choose a reason for hiding this comment

DavidLandup0 Aug 23, 2023

Choose a reason for hiding this comment

ianstenbit Aug 23, 2023

Choose a reason for hiding this comment

ianstenbit commented Aug 21, 2023

ianstenbit commented Aug 21, 2023

DavidLandup0 commented Aug 22, 2023

ianstenbit left a comment

Choose a reason for hiding this comment

ianstenbit commented Aug 24, 2023

DavidLandup0 commented Aug 24, 2023

ianstenbit commented Aug 24, 2023

ianstenbit commented Aug 24, 2023

DavidLandup0 commented Aug 24, 2023

ianstenbit commented Aug 24, 2023 • edited Loading

DavidLandup0 commented Aug 24, 2023

DavidLandup0 commented Jul 13, 2023 •

edited

Loading

DavidLandup0 commented Jul 17, 2023 •

edited

Loading

DavidLandup0 commented Jul 23, 2023 •

edited

Loading

ianstenbit commented Aug 24, 2023 •

edited

Loading