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

Add BASNet to keras hub #1984

Merged
merged 18 commits into from
Dec 19, 2024
Merged

Conversation

laxmareddyp
Copy link
Collaborator

@laxmareddyp laxmareddyp commented Nov 12, 2024

Refer this issue for more details

Please refer working demo notebook

@laxmareddyp laxmareddyp changed the title Port BASNet to keras hub Add BASNet to keras hub Nov 12, 2024
@laxmareddyp laxmareddyp added the kokoro:force-run Runs Tests on GPU label Nov 12, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Nov 12, 2024
Copy link
Collaborator

@divyashreepathihalli divyashreepathihalli left a comment

Choose a reason for hiding this comment

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

Thanks for the PR Laxma! it looks great! can you add a demo colab to verify outputs? Thanks!

keras_hub/src/models/basnet/basnet_test.py Show resolved Hide resolved
Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

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

Started a review. But I think this might need some bigger considerations.

Some questions:

  • What are the weight sources we are using here?
  • What are the outputs? Looks like there is a list of outputs instead of a normal classification style output.
  • How does training this model look?

In general, I think there there will be a lot of core infrastructure that will break if the arch does not have a backbone class. DeepLabV3 is probably much closer to what we will need to go for.

But I think I need answers to the above questions to better suggest a design.

keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Show resolved Hide resolved
@laxmareddyp
Copy link
Collaborator Author

laxmareddyp commented Nov 27, 2024

Some questions:

  • What are the weight sources we are using here?

The weights conversion script uses the weights from https://keras.io/examples/vision/basnet_segmentation/ . There's also a Keras BASNet model available on Kaggle, but it failed to load for me so far, and I don't know if where it comes from, if it's better or if it's the same model anyway.

  • What are the outputs? Looks like there is a list of outputs instead of a normal classification style output?

Yes, the direct model output is a list for refinement and prediction modules. The different stages of the model are trained with the same segmentation mask data. For inference, access to different stages of prediction modules would not be required for the user, but it might be interesting and doesn't hurt.

  • How does training this model look?

Training is simply performed on segmentation map data. For reference check this keras-cv notebook published in keras-io examples and issue

@mattdangerw
Copy link
Member

Maybe for now just directly subclass resnet without changes as BasnetBackbone? That way we can satisfy all the infra that will expect every new arch to have a backbone class.

The broader question I have is how to make our ImageSegmenter tasks interchangeable in terms of data in and prediction out. The goal of the task API is to allow users to swap different arches while keep the overall training/inference code the same. But it seems like there's a number of different ways to pose the segmentation problem. Not sure what to do here, but we can at least add a backbone class for now.

@laxmareddyp
Copy link
Collaborator Author

Maybe for now just directly subclass resnet without changes as BasnetBackbone? That way we can satisfy all the infra that will expect every new arch to have a backbone class.

The broader question I have is how to make our ImageSegmenter tasks interchangeable in terms of data in and prediction out. The goal of the task API is to allow users to swap different arches while keep the overall training/inference code the same. But it seems like there's a number of different ways to pose the segmentation problem. Not sure what to do here, but we can at least add a backbone class for now.

I have created a BASNetBackbone and moved the architecture implementation there (similar to how it is done in DeepLab). BASNetImageSegmenter now outputs a single segmentation map, therefore should be compatible to your described interface. I've implemented a compute_loss() to train BASNet's prediction modules as well.

@mattdangerw
Copy link
Member

Nice! Will take a look. Test failures look to be legit and basnet related, please take a look!

Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

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

Let's try this out! I wonder if we will need to make more tweaks, having the num classes on the backbone itself is kinda funky, but do get why we are doing it for this architecture.

Left a number of style comments.

keras_hub/src/models/basnet/__init__.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet_backbone.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet_backbone.py Outdated Show resolved Hide resolved
keras_hub/src/models/basnet/basnet_presets.py Outdated Show resolved Hide resolved
@mattdangerw
Copy link
Member

@laxmareddyp check the test failures! ~6 lines are too long and need fixing

@laxmareddyp
Copy link
Collaborator Author

@laxmareddyp check the test failures! ~6 lines are too long and need fixing

@mattdangerw I am modifying the code related to doc strings, other minor suggestions you mentioned and will push the changes and hope it will not show up this time. Thanks

@laxmareddyp
Copy link
Collaborator Author

Let's try this out! I wonder if we will need to make more tweaks, having the num classes on the backbone itself is kinda funky, but do get why we are doing it for this architecture.

num_classes: Unlike DeepLab, in BASNet this parameter is quite involved in the architecture. This is because BASNet has multiple model outputs, Therefore, we cannot construct the task by just adding a conv layer with num_classes outputs to the backbone. A possibility to achieve this goal would be splitting the backbone in several parts, which are used individually in the task, but the effort for this would (in my opinion) outweigh the benefits by far.

@laxmareddyp laxmareddyp added the kokoro:force-run Runs Tests on GPU label Dec 17, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Dec 17, 2024
Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

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

Looks good overall. I'll drop the conversion script as I merge. Let's do that separately, need some work

tools/checkpoint_conversion/convert_basnet_checkpoints.py Outdated Show resolved Hide resolved
tools/checkpoint_conversion/convert_basnet_checkpoints.py Outdated Show resolved Hide resolved
tools/checkpoint_conversion/convert_basnet_checkpoints.py Outdated Show resolved Hide resolved
tools/checkpoint_conversion/convert_basnet_checkpoints.py Outdated Show resolved Hide resolved
@laxmareddyp laxmareddyp added the kokoro:force-run Runs Tests on GPU label Dec 17, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Dec 17, 2024
Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

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

lgtm!

@mattdangerw mattdangerw merged commit 7dc39fc into keras-team:master Dec 19, 2024
10 checks passed
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.

4 participants