The Keras package is an open-source library that provides a Python interface for deep learning. Keras is intended to be a user-friendly, modular, and extensible way to enable fast experimentation with deep neural networks.
Keras has an approachable API that originally supported multiple backends including TensorFlow, Microsoft Cognitive Toolkit, Theano, and PlaidML. With the introduction of Keras version 2.4, however, only the TensorFlow backend was supported.
With Keras version 3, however, the package provides APIs for using three backends: TensorFlow, JAX, and PyTorch. This is a nice change in our view, as PyTorch has gained traction for deep learning research and general use in many labs.
Technical Note: as of the date of publication of this article, the current version of Keras is 3.2.1. We recommend using the latest releases, as they are still working out some small bugs and API improvements.
Why would you want to use a different backend?
Some machine learning backend frameworks can be more optimized for specific tasks or hardware. For smaller, fully connected models that are common in Earth Systems Science research, this might not matter as much. But some of the speed-ups reported with using JAX might be worth looking into for larger problems. This article about JAX covers some of the advantages: Why You Should (or Shouldn't) be Using Google's JAX in 2023.
One reason to use the PyTorch ecosystem is the Fully Sharded Data Parallel API which allows for massive scaling on the largest of clusters.
While there are pros and cons to each backend, there is not one “best” option — especially when dealing with large, complex problems. Keras 3 benchmarks provides some additional insights into the suitability of different backend frameworks for different situations.
Is TensorFlow Dead?
Just like Fortran and Matlab, TensorFlow will be around for a long time. But with JAX having official support from Google, TensorFlow may not get the development support it has in the past. If you have a current project that uses TensorFlow, I would not worry about it today, but this evolving landscape is something to keep in mind for the future. (That future may be a long way off!) If problems with support for TensorFlow do start to emerge, with Keras 3 it should not be a major change to switch backends.
Why Does Unidata Care?
For deep learning training I (Thomas) will be using a Keras 3 API exclusively. It more closely resembles the scikit-learn api and I find it to be easier to explain. Other people might feel differently, but I currently do not see the downsides to teaching the new Keras api. This API is also used for some JupyterHubs already at Universities.
Unidata’s Science Gateway Now Offers Access to Jupyter Hubs with Keras 3 loaded!
Please get in touch with our Science Gateway team at support-gateway@unidata.ucar.edu to inquire about access to a GPU enabled JupyterHub instance with Keras 3! While we offer this access free of charge to our community, access to these resources may be limited depending on demand.
An example notebook illustrating use of the Keras 3 API is here: Modified Weather Bench Code
You can find more information about Keras 3 in Introducing Keras 3.0.
Thomas Martin is an AI/ML Software Engineer at the NSF Unidata Program Center. Have questions? Contact support-ml@unidata.ucar.edu or book an office hours meeting with Thomas on his Calendar.
