Suppress training output warnings and remove intermediate epoch PNG visualizations in Focal Modulation Network tutorial

[LakshmiKalaKadali]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello, 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 significantly updates the Focal Modulation Network tutorial to be fully compatible with Keras 3, ensuring backend-agnostic operation across JAX, PyTorch, and TensorFlow. The changes involve a comprehensive refactor of the data loading pipeline using PyDataset, modernization of model layer implementations, and integration of preprocessing directly within the model. Additionally, the tutorial's output has been cleaned up by suppressing verbose warnings and removing redundant intermediate visualization images, enhancing clarity and user experience.

Highlights

• Keras 3 Migration: The tutorial has been updated to Keras 3, making the model backend-agnostic by utilizing keras.ops for mathematical operations and keras.utils.PyDataset for data loading. This involves significant refactoring of imports, layer definitions, and data pipelines.
• Warning and Logging Suppression: Added explicit code to suppress TensorFlow C++ logging and various Keras/TensorFlow warnings, leading to a cleaner and less verbose training output.
• Streamlined Model Architecture: The FocalModulationNetwork and its sub-layers (PatchEmbed, MLP, FocalModulationLayer, FocalModulationBlock, BasicLayer) have been refactored to remove redundant attributes, simplify __init__ and call signatures, and integrate data preprocessing directly into the model's call method for better backend compatibility.
• Removed Intermediate Visualizations: The tutorial no longer generates and displays intermediate PNG visualizations of gates and modulators during training epochs, simplifying the output and focusing on the final results.
• Updated Tutorial Metadata and Configuration: The 'Last modified' date has been updated, and the tutorial's global configuration parameters like BATCH_SIZE and EPOCHS have been adjusted for potentially faster execution or better demonstration with Keras 3.

🧠 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/focal_modulation_network.py
  • Updated 'Last modified' date and added Keras 3 conversion attribution.
  • Replaced TensorFlow-specific imports with Keras 3 equivalent imports (keras, keras.ops).
  • Implemented environment variable settings to suppress TensorFlow C++ logging and Python warnings.
  • Refactored global configuration parameters, including BATCH_SIZE and EPOCHS.
  • Replaced tf.data pipeline with a custom FocalDataset class inheriting from keras.utils.PyDataset for backend-agnostic data handling.
  • Simplified PatchEmbed layer by removing redundant attributes and updating call method to use ops.shape.
  • Updated MLP function to use string literal 'gelu' for activation and removed type hints.
  • Corrected GeLU activation link in markdown comments.
  • Refactored FocalModulationLayer to simplify initialization, use ops.split and ops.softmax, and streamline context aggregation.
  • Simplified FocalModulationBlock and BasicLayer initializations and call methods, removing unused parameters and updating tensor operations to ops.
  • Integrated image preprocessing (rescaling, resizing, cropping, flipping) directly into the FocalModulationNetwork's call method.
  • Removed pos_drop and flatten layers from FocalModulationNetwork.
  • Modified display_grid function to use ops.convert_to_numpy and ops.norm for backend-agnostic visualization.
  • Updated TrainMonitor callback to use test_batch_images and integrate upsampler as a class attribute, removing intermediate image generation.
  • Adjusted WarmUpCosine learning rate schedule to use ops.cast and ops.where for backend compatibility.
  • Updated model compilation and fitting to use the new model instance and removed TrainMonitor callback from the fit call to prevent intermediate image generation.
  • Modified plotting code to create two subplots for loss and accuracy.
• examples/vision/ipynb/focal_modulation_network.ipynb
  • Updated 'Last modified' date in notebook metadata.
  • Reflected all code changes from focal_modulation_network.py within the notebook cells.
  • Removed output cells containing intermediate training progress and image visualizations.
• examples/vision/md/focal_modulation_network.md
  • Updated 'Last modified' date in the markdown header.
  • Reflected all code changes from focal_modulation_network.py within the markdown code blocks.
  • Removed image paths for intermediate training visualizations.

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 updates the Focal Modulation Network tutorial to be compatible with Keras 3. The changes include migrating from TensorFlow-specific APIs to the backend-agnostic keras.ops, replacing the tf.data pipeline with keras.utils.PyDataset, and moving data augmentation steps inside the model. Additionally, it suppresses training output warnings and removes intermediate epoch visualizations to provide a cleaner execution experience. The code has been simplified and modernized, but I've identified a potential issue in the FocalModulationLayer's call method where ops.split is used with incorrect arguments, which could lead to a runtime error.