Unleashing the Power of Google Colab, TensorFlow and Keras for AI Enthusiasts

In an era where the pace of technological advancement accelerates daily, the collaborative synergy of TensorFlow, Keras, and Google Colab represents a beacon of innovation in artificial intelligence (AI) and machine learning (ML). This powerful trio not only equips developers with sophisticated tools but also heralds a new age of collaborative development. It’s a narrative of transformation, where the complexities of AI are unraveled, making way for a future where anyone with an idea can contribute to the technological frontier.

We will explore various tips, tricks, and uncommon usage scenarios specifically tailored for experienced TensorFlow developers using Google Colab. These advanced techniques will help you streamline your development process, leverage hardware acceleration, implement distributed training, and utilize the Google Cloud Platform. Additionally, we will provide useful resources that you can refer to for further learning.

Whether you are a seasoned TensorFlow developer or just beginning to explore its advanced capabilities, this post will equip you with valuable insights and practices that will elevate your TensorFlow development experience with Google Colab.

Enter Google Colab

Google Colab stands out not merely as a tool but as a collaborative platform that democratizes access to AI development. Its provision of free computational resources, such as GPUs and TPUs, eliminates traditional barriers to entry, allowing for a broader range of innovators to participate in AI development. The platform’s real-time collaboration feature fosters a community where ideas are shared freely, accelerating the iterative process of innovation. This environment is crucial for nurturing groundbreaking AI solutions, exemplified by the ability for multiple users to interact with the same Jupyter Notebook simultaneously, mirroring the collaborative essence of modern scientific exploration.

Free GPU and TPU resources

Colab also comes with free access to GPUs and TPUs. Here you can carry out computationally demanding activities without the need for expensive hardware. Deep learning models or more so LLMs (Large Language Models) recently, which are constantly hungry for pricey computational resources, can notably benefit from this.

Collaborative environment

Thanks to its integration with Google Drive, Google Colab enables you to work in real-time with your colleagues and collaborators. This feature makes it possible for several people to see and edit the same notebook at once, facilitating efficient teamwork.

Pre-installed libraries

Colab makes it simple for you to get started on your work right away by pre-installing a number of well-known libraries and packages. You can also add any other packages you require to increase its adaptability.

AI-powered Coding (just in!)

As Google Colab rolls out these AI coding features, we can expect a significant impact on the world of programming. With improved speed, quality, and comprehension, Python developers will be able to produce better code in less time. Keep an eye out for these exciting new features if you happen to be a paid subscriber.

TensorFlow and Keras: Bridging Complexity and Creativity

The integration of TensorFlow and Keras signifies a pivotal shift towards more accessible and flexible AI development. TensorFlow’s adaptability, paired with Keras’ intuitive interface, offers a dual advantage — sophistication without the steep learning curve. This combination encourages a broader spectrum of developers to engage with deep learning, pushing the boundaries of what’s possible. Real-world applications, such as TensorFlow’s role in enhancing Saatchi Art’s “8th Wall” and the immersive experience of IKEA’s “IKEA Place” app, underscore the practical impact of these technologies. These examples not only serve as inspiration but also as concrete evidence of how TensorFlow and Keras facilitate the creation of applications that were once deemed futuristic.

Colab subscription tiers (as of 2023)

You can choose the right plan for your project here or just try out the free one if you’re just getting started.

+---------+----------------+----------------------+------------------------+
|         | Google Colab   | Google Colab Pro     | Google Colab Pro+      |
+---------+----------------+----------------------+------------------------+
| GPUs    | K80, T4        | T4, V100, A100       |  T4, V100, A100        |           
+---------+----------------+----------------------+------------------------+
| RAM     | Mostly 12GB    | Standard/Premium     | Much higher limits     |     
+---------+----------------+----------------------+------------------------+
| Price   | Free           | $9.99/month          | $49.99/month           |
+---------+----------------+----------------------+------------------------+

(RAM, GPUs and TPUs here are always subject to change)

Harnessing the Power of Google Colab and TensorFlow

Now let us get right away into combining the strengths of Google Colab and TensorFlow and open up a world of possibilities for our AI enthusiasts. As most of you might be well aware of some of the simple ways to leverage TensorFlow to build your models and ML workflows, we will instead mostly focus on utilizing TensorFlow and Colab in unusual yet proven-to-be-useful ways. Diving deeper, this article explores advanced techniques that leverage the unique strengths of Google Colab and TensorFlow. From notebook tricks that enhance the development experience to integrating TensorFlow with TPU backends for accelerated performance, the focus is on practical, proven strategies that can significantly enhance AI projects.

Notebook tricks

As you may know for various Python to Javascript and Javascript to Python interactions, APIs are provided; however, they are often only accessible for outputs that have been executed within the current session.

In an effort to provide a rich viewing experience, Colaboratory stores all of the resources required to see outputs in the notebook itself, where they can be accessed later on even if the notebook is not connected to a kernel.

display.Javascript to execute JavaScript from Python

Evaluate a Javascript expression from Python with eval_js

Programmatically force restart the Colab kernel

Restarting the runtime in Google Colab is sometimes necessary to ensure that the changes made to the notebook environment take effect. There are several reasons why you might need to restart the runtime, including:

• When you install or upgrade a package using !pip install, it may be necessary to restart the runtime for the new package version to be correctly loaded and used in your notebook.
• Restarting the runtime can help free up memory if your notebook has been running for a long time or if you’ve been experimenting with large datasets and models.
• Sometimes, you may need to reset the state of your notebook to start from a clean slate or to avoid conflicts arising from variables being overwritten or modified.
• Occasionally, you might encounter errors or issues that are difficult to resolve without restarting the runtime.

We will be performing the task programmatically to ensure a seamless experience when running the notebook. This approach eliminates the need for the user to manually navigate to the Runtime menu, streamlining the process and improving efficiency.

Browsing to servers executing on the kernel

This will create an iframe browsing the HTTP server hosted on the machine your kernel is running on.

Alternatively to view the server in a separate browser tab

Running Machine Learning applications

Similarly a well known Python package known as Gradio leverages the same aspect of browsing to a server running inside the Colab VM by making a Machine Learning app publicly accesible as you can see below.

Import your data from Google Drive into Colab

Since Google Drive’s storage is way slower than that of the Colab VM, you can expect slow read/write speeds after mounting the drive folder in Colab. This is because you are accessing files through Google Drive, which has a longer access time. One way to improve this would be to copy the files from your Drive to the Colab instance. To do this you may rely on the PyDrive APIs available in Colab.

To simplify things, utilizing the fantastic gdown utility is another option. It should be noted that even if gdown is already installed on Colab instances, you still need to upgrade the package for it to function. You’ll need the file ID, which you can find by selecting Right-Click → Get link → Anyone with link on the file in Drive, and then copying the ID (given in bold below) from a supplied GDrive URL.

Although moving all of your data to Google Cloud Storage and letting your Colab instance access it that way is the most advised way to access your data in Colab and perform any type of preprocessing or training on it.

Installing a specific version of TensorFlow for your TPU backend

The following script can come in handy if you ever feel like switching to a specific TensorFlow version.

Data preprocessing, visualization and accessibility

Use Colab’s preinstalled libraries like pandas, NumPy, matplotlib and Colab’s own features to preprocess, visualize data and maximize accessibility, ensuring your data is ready for ML models to use.

Forms

Forms offer a straightforward method for parameterizing your code. Any changes you make to the form’s value will automatically update the corresponding value in your code.

From a code cell, select Insert → Add form field.

To explore various use-case scenarios, I encourage you to experiment with Colab’s official notebook for using forms.

Jupyter/Colab Widgets

Jupyter Widgets can be used to add interactivity to notebooks. Ipywidgets are interactive widgets for Jupyter Notebooks, which allow you to create dynamic user interfaces for your data analysis and visualization projects. With Ipywidgets, you can create sliders, buttons, text inputs, and other interactive elements, which update in real-time as you manipulate them. Read more on how you can leverage these widgets in this notebook.

Colab comes with its own widgets which you can take a look here.

Embed media

The following Python code is designed to convert a sequence of images (from a sample video) into a GIF and embed it into an IPython notebook, specifically for display purposes in a Jupyter notebook or similar IPython environments (Colab in our case).

Source: https://github.com/tensorflow/docs/blob/master/tools/tensorflow_docs/vis/embed.py

• The embed_data function is defined to embed the data as an HTML tag with a data-url, converting data to base64 and wrapping it in the appropriate HTML tag for images and videos. It raises an error if the data is not an image or video.
• The embed_file function is defined to read the bytes from a specified file path, guess the mimetype of the file, and then pass this data to the embed_data function. The function returns the resulting HTML.
• The to_gif function is defined to convert a numpy array of images into a GIF file. It scales the image data to an 8-bit range, writes the images to a GIF file using the imageio library, and then uses the embed_file function to embed the resulting GIF into the IPython notebook display.

Finally, the to_gif function is called with sample_video as an argument to execute the above processes on the sample_video data. This code is a useful way to visualize sequences of images or videos directly within a Colab notebook.

Magic commands

If you ever wanted to overwrite a file or prefer using a code cell to write some content to a file you can use the %writefile magic command to do so.

Magics like these make it easy to develop, test, compare, and evaluate from within a Colab notebook. There are several other magic commands at your disposal too. You can find the entire list here.

Making your Colab notebooks more accessible in GitHub

Colab can load public github notebooks directly, with no required authorization step.

For example, consider the notebook at this address: https://github.com/googlecolab/colabtools/blob/master/notebooks/colab-github-demo.ipynb.

The direct colab link to this notebook is: https://colab.research.google.com/github/googlecolab/colabtools/blob/master/notebooks/colab-github-demo.ipynb.

Give more life to your dataframes

Pandas is a wonderful tool for analyzing data in Python but wouldn’t it be nice if your dataframes could be a bit more interactive? Google Colab offers an extension that further enhances these dataframes, allowing for dynamic filtering, sorting, and exploration.

Would you like to enable this interactive data table display? If so, you can simply execute the following in Colab

Try it out yourself here.

TensorFlow Datasets (tf.data) and Google Cloud Storage

Some datasets are available to load directly from the public TFDS GCS bucket (gs://tfds-data/datasets/)without any authentication:

• If tfds.load(..., try_gcs=False) (default), the dataset will be copied locally in ~/tensorflow_datasets during download_and_prepare.
• If tfds.load(..., try_gcs=True), the dataset will be streamed directly from GCS (download_and_prepare will be skipped).

You can check whether a dataset is hosted on the public bucket with tfds.is_dataset_on_gcs('mnist').

And finally load it using tfds.load.

This is especially useful when you want to load really big datasets. (as of right now there aren’t many datasets in GCS, but it goes to show that you generally should be hosting your datasets on GCS so that Colab can use GCS to efficiently load data and train your models.

Efficient data loading using TFRecords

As you already know, hardware accelerators are very fast. The challenge is often to feed them data fast enough to keep them busy. Google Cloud Storage (GCS) is capable of sustaining very high throughput but as with all cloud storage systems, initiating a connection costs some network back and forth. Therefore, having our data stored as thousands of individual files is not ideal. We are going to batch them in a smaller number of files and use the power of tf.data.Dataset to read from multiple files in parallel.

The TFRecord file format

The TFRecord file format, based on protobuf, is TensorFlow’s preferred method for storing data. While other serialization formats are also compatible, datasets can be directly loaded from TFRecord files for convenience.

Model building, saving and training

Saving/Loading models with TensorFlow on TPU

When training on remote devices, such as Cloud TPUs, and saving and loading from a local I/O device, you need to utilize the experimental_io_deviceoption within tf.saved_model.SaveOptionsand tf.saved_model.LoadOptions. This is to specify that the I/O device is set to localhost.

Leverage TensorFlow’s Keras API to build and train deep learning models easily. Colab’s GPU and TPU support significantly reduce training time for large models. Now let’s look at effectively making use of these free accelerators at our disposal.

GPUs

Here are some performance tips when using mixed precision on GPUs.

• Increasing your batch size: If it doesn’t affect model quality, try running with double the batch size when using mixed precision. As float16 tensors use half the memory, this often allows you to double your batch size without running out of memory. Increasing batch size typically increases training throughput, i.e. the training elements per second your model can run on.
• Ensuring GPU Tensor Cores are used: As mentioned previously, modern NVIDIA GPUs use a special hardware unit called Tensor Cores that can multiply float16 matrices very quickly. However, Tensor Cores requires certain dimensions of tensors to be a multiple of 8.

Check the NVIDIA deep learning performance guide for optimizing Tensor Cores.

TPUs

As with GPUs, you should try doubling your batch size when using Cloud TPUs because bfloat16tensors use half the memory. Doubling batch size may increase training throughput.

TPUs do not require any other mixed precision-specific tuning to get optimal performance. They already require the use of XLA. TPUs benefit from having certain dimensions being multiples of , but this applies equally to the float32type as it does for mixed precision.

Check the Cloud TPU performance guide for general TPU performance tips, which apply to mixed precision as well as float32 tensors.

XLA

Under the hood, a compiler called XLA (accelerated Linear Algebra compiler) transforms the TensorFlow graph of computation nodes into TPU machine code. This compiler also performs many advanced optimizations on your code and your memory layout.

To read more about XLA and tf.function, check out TensorFlow’s official tutorial to train a TensorFlow model to classify the MNIST dataset, where the training function is compiled using XLA.

Addressing Data Bottlenecks, Data Types and Compatibility Constraints

• TPUs are very fast and ingesting data often becomes the bottleneck when running on them. There are tools you can use to detect data bottlenecks and other performance tips in the TPU Performance Guide.
• int8 or int16 numbers are treated as int32. The TPU does not have integer hardware operating on less than 32 bits.
• Some TensorFlow operations are not supported. The good news is that this limitation only applies to training code i.e. the forward and backward pass through your model. You can still use all Tensorflow operations in your data input pipeline as it will be executed on CPU.
• tf.py_func is not supported on TPU.

General tips while porting a notebook to use TPUs

TPU Initialization

Initialize your TPUs at the start of your Colab notebook.

Convert your models to use the TPU strategy.

Build and compile your model within the strategy scope.

Utilize bfloat16

TPUs are designed to run best with bfloat16 (Brain Floating Point) data type. While the reduced precision as compared to float32 might seem like a disadvantage, in practice, neural networks often perform just as well. You might see a considerable speed-up by converting your model's computations to bfloat16.

Model compilation

tf.function is a tool used to convert Python code into Python-independent dataflow graphs, enhancing the performance and portability of your models. This transformation is essential for utilizing SavedModel.

TPUs perform best when running large chunks of computation at once. You can use tf.function to compile your critical code sections into a single computation graph.

In order to compile the computation using XLA, place it inside tf.function with jit_compile=True.

Gradient Clipping

In some cases, gradient clipping can help improve the performance of a model on a TPU by preventing the gradients from becoming too large and causing numerical stability problems.

Model deployment and serving

Deploy and serve your trained models across different platforms using TensorFlow Serving, which supports environments ranging from local servers to the cloud. For deployment on mobile devices and web applications, use TensorFlow Lite and TensorFlow.js, respectively.

Exporting to a SavedModel for TF Serving

To deploy a model using TensorFlow Serving:

1. Train a neural network on the MNIST dataset using TensorFlow.
2. Export the trained model as a SavedModel.
3. Serve the model with TensorFlow Serving.

Finally, we use requests to communicate with the TensorFlow Serving server for inference.

You can follow the steps in TensorFlow’s official Colab notebook here to try it out for yourself.

Serving models with Vertex AI

You can use TensorFlow Serving functions with Vertex AI Prediction and Vertex AI Endpoints to deploy your models in Google Cloud. The way you do this is pretty simple.

After deploying your Model resource to an Endpoint resource, you can execute online predictions by submitting prediction requests directly to that Endpoint resource.

For a complete end-to-end guide, it’s best to check out Google Cloud Platform’s official notebook. Go ahead and give it a try by importing the notebook in Vertex AI Workbench.

Accelerated inference on NVIDIA GPUs using TensorRT

The NVIDIA TensorRT is a C++ library that facilitates high performance inference on NVIDIA graphics processing units (GPUs). TensorRT takes a trained network, which consists of a network definition and a set of trained parameters, and produces a highly optimized runtime engine which performs inference for that network.

The TensorFlow integration with TensorRT (TF-TRT) optimizes and executes compatible subgraphs, allowing TensorFlow to execute the remaining graph. While you can still use TensorFlow’s wide and flexible feature set, TensorRT will parse the model and apply optimizations to the portions of the graph wherever possible.

Machine learning reproducibility

Machine learning reproducibility is critical for ensuring scientific validity, facilitating effective debugging and model improvement, enabling accurate comparisons between methods, fostering efficient knowledge transfer, and avoiding redundant work. It plays an integral role in advancing the field of machine learning. With all that said, if you want to have more stable results across different training runs, always remember to fix your seeds and set any determinism flags available in TensorFlow.

Experiment tracking and collaboration

Document your experiments, track model performance, and share your work using TensorBoard and Google Drive, making it easy to collaborate with peers and showcase your findings.

The TensorBoard-Colab integration

Effective experiment tracking and visualization are crucial for understanding the performance of machine learning models and making informed decisions during the development process. Google Colab integrates well with popular experiment tracking tools like TensorBoard and Neptune.

Colab also provides a magic command to use Tensorboard directly from the notebook. You just need to set the logs directory location

You can learn to use TensorBoard in Colab by following the steps in the official notebook.

This narrative is more than a guide; it’s a journey through the transformative power of AI tools. It illustrates a path from the initial spark of an idea to its realization as a solution that reshapes industries. Through the lens of TensorFlow, Keras, and Google Colab, we see a world where AI is not just a tool for developers but a catalyst for innovation across sectors, making what once seemed impossible, possible.

As we stand at the precipice of a new era in AI development, the combined force of TensorFlow, Keras, and Google Colab offers a beacon of hope and innovation. This article serves as both a testament to their potential and a guide for harnessing their power. It encourages developers to explore these advanced capabilities, ensuring they are well-equipped to contribute to the ever-evolving landscape of AI. With these tools at their disposal, developers are not just participants in the AI revolution; they are its architects, shaping the future of technology.

Implementing these strategies in your ML workflows should help you as a developer in streamlining your projects, improving teamwork effectiveness, and maximizing TensorFlow’s potential with Google Colab.