Migration of masked_image_modeling_with_autoencoders to keras 3

[maitry63]

This PR migrates the Masked image modeling with Autoencoders tutorial from Keras 2 to Keras 3. The implementation is now fully backend-agnostic, allowing it to run seamlessly on JAX, PyTorch, and TensorFlow (notebook).

[gemini-code-assist]

Summary of Changes

Hello @maitry63, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request focuses on updating the 'Masked image modeling with Autoencoders' tutorial to be fully compatible with Keras 3. The primary goal was to achieve backend-agnosticism, allowing the tutorial to run seamlessly across various Keras 3 backends like JAX, PyTorch, and TensorFlow. This involved a comprehensive refactoring of tensor operations to use keras.ops, modernizing the data loading pipeline, and adapting custom training logic to the new Keras 3 API.

Highlights

• Keras 3 Migration: The tutorial has been fully migrated from Keras 2 to Keras 3, ensuring compatibility with the latest Keras API and features.
• Backend Agnosticism: All tensor operations now utilize keras.ops, making the implementation backend-agnostic and runnable on JAX, PyTorch, and TensorFlow.
• Data Pipeline Refactoring: The data loading mechanism has been updated to use keras.utils.PyDataset, replacing tf.data.Dataset.from_tensor_slices for improved Keras 3 integration.
• Custom Training Loop Updates: The train_step and test_step methods within the custom MAE model have been revised to align with Keras 3's training loop conventions, including simplified gradient application and metric reporting.
• Hyperparameter Adjustments: The BATCH_SIZE and EPOCHS for pretraining and linear probing have been reduced to facilitate quicker execution and demonstration within the tutorial context.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• examples/vision/masked_image_modeling.py
  • Updated metadata including 'Last modified' date and 'Converted to Keras 3 by' author.
  • Replaced direct tensorflow imports and operations with keras.ops for backend-agnostic tensor manipulation.
  • Introduced keras.utils.PyDataset for data loading, replacing tf.data.Dataset.from_tensor_slices.
  • Modified Patches layer to use ops.image.extract_patches and ops.reshape.
  • Updated show_patched_image and reconstruct_from_patch methods to use ops.convert_to_numpy and ops functions for image manipulation.
  • Added a docstring to the PatchEncoder class for better clarity.
  • Refactored PatchEncoder initialization and call method to use keras.Variable, keras.random.normal, and various keras.ops functions for random indices, tiling, and gathering.
  • Changed activation function tf.nn.gelu to string literal 'gelu' in mlp function.
  • Removed epsilon parameter from LayerNormalization and dropout from MultiHeadAttention where defaults are sufficient or handled differently in Keras 3.
  • Updated calculate_loss to use ops.concatenate.
  • Revised train_step and test_step to accept data parameter and simplified gradient application using self.trainable_variables.
  • Modified WarmUpCosineDecay to use ops.convert_to_tensor, ops.cast, ops.cos, and ops.where for learning rate scheduling.
  • Adjusted mae_model.evaluate output handling to accommodate Keras 3's potentially dictionary-based metrics.
  • Reduced BATCH_SIZE to 32, EPOCHS to 1 for pretraining, and linear_probe_epochs to 2 for faster tutorial execution.

Activity

• The pull request was created by maitry63 with the explicit goal of migrating the 'Masked image modeling with Autoencoders' tutorial to Keras 3 and ensuring backend-agnosticism. All changes reflect this migration effort.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request successfully migrates the Masked Image Modeling example to Keras 3, with a focus on making it backend-agnostic. The changes include updating TensorFlow-specific APIs to their keras.ops equivalents, introducing a PyDataset for data loading, and simplifying parts of the model implementation. However, there are several critical issues that prevent the code from being truly backend-agnostic as claimed, most notably the use of tf.GradientTape in the custom train_step and tf.gather. I've also pointed out some inconsistencies in the model architecture and opportunities for code cleanup and simplification.

[maitry63]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request does a great job of migrating the Masked Image Modeling example to Keras 3, making it backend-agnostic. The changes correctly replace TensorFlow-specific APIs with keras.ops and adopt Keras 3 conventions like using keras.utils.PyDataset. I've found one critical issue with the handling of the return value from model.evaluate() that will cause the script to fail. Additionally, there's a small oversight in the create_encoder function that was missed during the migration. In create_encoder, the MultiHeadAttention layer is still initialized with a dropout argument. This argument has been removed in Keras 3, and its presence will cause an error. It should be removed, just as it was in create_decoder. My detailed comment provides a suggestion for the other point.

[maitry63]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request successfully migrates the Masked Image Modeling example to Keras 3, making most of the code backend-agnostic by replacing TensorFlow-specific APIs with keras.ops. The changes are generally well-executed, improving code clarity and adhering to Keras 3 best practices. However, I've identified a few critical issues that need to be addressed. There's a bug in how the model evaluation results are handled, a regression in the data pipeline that removes shuffling for the training set, and a redundant image resizing step that alters the augmentation logic. Additionally, the linear probing part of the example still relies on tf.data.Dataset, which prevents it from being truly backend-agnostic. Addressing these points will ensure the example is correct, efficient, and fully compatible with all Keras backends.

[maitry63]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request successfully migrates the Masked Image Modeling example to Keras 3, making it backend-agnostic. The changes correctly adopt keras.ops for operations and keras.utils.PyDataset for data handling, and the custom training loop has been effectively refactored. However, I've identified a critical issue with a duplicated class definition that needs to be resolved. Additionally, there are several opportunities for cleanup, such as removing unused imports and commented-out code, which will enhance the overall quality and readability of the example. I've left specific comments with suggestions for these improvements.

[sachinprasadhs]

Thanks, made few comments, please address it.

[sachinprasadhs]

Few more comments to address.

[sachinprasadhs]

Thanks, Looks good. Please generate .md and .ipynb files.

[github-actions]

This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

[sachinprasadhs]

LGTM