TensorFlow Lite vs PyTorch Mobile for On-Device Machine Learning

In the recent world of technology development and machine learning it’s no longer confined in the micro cloud but in mobile devices. As we know, TensorFlow Lite and PyTorch Mobile are two of the most commercially available tools for deploying models directly on phones and tablets. TensorFlow Lite and PyTorch mobile, both, are developed to operate on mobile, yet they stand distinct in their pros and cons. Here in this article we are to know what TensorFlow Lite is, what is PyTorch Mobile, their applications and differences between both.

Learning Outcomes

• Overview of device machine learning and why it is beneficial rather than cloud based systems.
• Learn about TensorFlow Lite and PyTorch Mobile used for mobile application deployment.
• How to convert trained models for deployment using TensorFlow Lite and PyTorch Mobile.
• Compare the performance, ease of use, and platform compatibility of TensorFlow Lite and PyTorch Mobile.
• Implement real-world examples of on-device machine learning using TensorFlow Lite and PyTorch Mobile.

This article was published as a part of the Data Science Blogathon.

What is On-Device Machine Learning?

We can perform AI on the mobile devices including smart phone, tablet or any other device using on device machine learning. We do not need to rely on services of clouds. These are fast response, security of sensitive information, and application can run with or without internet connectivity which are very vital in diverse applications; image recognition in real-time, machine translation, and augmented reality.

Exploring TensorFlow Lite

TensorFlow Lite is the TensorFlow version which is often used on devices with restricted capabilities. It works and is compatible with other operating systems such as the Android and the iPhone. It mainly centers itself in providing latency and high performance execution. As for TensorFlow Lite, there is a Model Optimizer that helps to apply certain methods, for example, quantization to models. This makes models faster and smaller for mobile deployment which is imperative in this practice to enhance efficiency.

Features of TensorFlow Lite

Below are some most important features of TensorFlow Lite:

• Small Binary Size: TensorFlow Lite binaries can be of very small size. It can be as small as 300KB.
• Hardware Acceleration: TFLite supports GPU and other hardware accelerators via delegates, such as Android’s NNAPI and iOS’s CoreML.
• Model Quantization: TFLite offers many different quantization methods to optimize performance and reduce model size without sacrificing too much accuracy.

PyTorch Mobile Implementation

PyTorch Mobile is the mobile extension of PyTorch. It is generally known for its flexibility in research and production. PyTorch Mobile makes it easy to take a trained model from a desktop environment and deploy it on mobile devices without much modification. It focuses more on the developer’s ease of use by supporting dynamic computation graphs and making debugging easier.

Features of PyTorch Mobile

Below are some important features of Pytorch Mobile:

• Pre-built Models: PyTorch Mobile provides a large range of pre-trained models that can be converted to run on mobile devices.
• Dynamic Graphs: It is one of PyTorch’s dynamic computation graphs that allow for flexibility during development.
• Custom Operators: PyTorch Mobile allows us to create custom operators, which can be useful for advanced use cases.

Performance Comparison: TensorFlow Lite vs PyTorch Mobile

When we discuss their performance, both frameworks are optimized for mobile devices, but TensorFlow Lite has high execution speed and resource efficiency.

• Execution Speed: TensorFlow Lite is generally faster due to its aggressive optimization, such as quantization and delegate-based acceleration. For example- NNAPI, and GPU.
• Binary Size: TensorFlow Lite has a smaller footprint, with binary sizes as low as 300KB for minimal builds. PyTorch Mobile binaries tend to be larger and require more fine-tuning for a lightweight deployment.

Ease of Use and Developer Experience

PyTorch Mobile is generally preferred by developers because of its flexibility and ease of debugging. It is because of dynamic computation graphs. This helps us to modify models at runtime, which is great for prototyping. On the other hand, TensorFlow Lite requires models to be converted to a static format before deployment, which can add complexity but result in more optimized models for mobile.

• Model Conversion: PyTorch Mobile allows us for direct export of PyTorch models, while TensorFlow Lite requires converting TensorFlow models using the TFLite Converter.
• Debugging: PyTorch’s dynamic graph makes it easier to debug models while they’re running, which is great for spotting issues quickly. With TensorFlow Lite’s static graph, debugging can be a bit difficult although TensorFlow provides tools such as Model Analyzer which can help us.

Supported Platforms and Device Compatibility

We can use both TensorFlow Lite and PyTorch Mobile on two major mobile platforms, Android and iOS.

TensorFlow Lite

When it comes to choosing which will support which hardware, TFLite is way more flexible. Due to the delegate system it supports not only CPUs and GPUs but also Digital Signal Processors (DSPs) and other chips that are deemed higher performers than the basic CPUs.

PyTorch Mobile

While PyTorch Mobile also supports CPUs and GPUs such as Metal for iOS and Vulkan for Android, it has fewer options for hardware acceleration beyond that. This means that TFLite may have the edge when we need broader hardware compatibility, especially for devices which have specialized processors.

Model Conversion: From Training to Deployment

The main difference between TensorFlow Lite and PyTorch Mobile is how models move from the training phase to being deployed on mobile devices.

TensorFlow Lite

If we want to deploy a TensorFlow model on mobile then it needs to be converted using the TFLite converter. This process can be optimized, such as quantization which will make the model fast and efficient for mobile Targets.

PyTorch Mobile

For PyTorch Mobile, we can save the model using TorchScript. The process is very simpler and easy, but it does not offer the same level of advanced optimization options that TFLite provides.

Use Cases for TensorFlow Lite and PyTorch Mobile

Explore the real-world applications of TensorFlow Lite and PyTorch Mobile, showcasing how these frameworks power intelligent solutions across diverse industries.

TensorFlow Lite

TFLite is a better platform for different applications that require quick responses such as real-time image classification or object detection. If we are working on devices with specialized hardware such as GPUs or Neural Processing Units. TFLite’s hardware acceleration features help the model run faster and more efficiently.

PyTorch Mobile

PyTorch Mobile is great for projects that are still evolving, such as research or prototype apps. Its flexibility makes it easy to experiment and iterate, which allows developers to make quick changes. PyTorch Mobile is ideal when we need to frequently experiment and deploy new models with minimal modifications.

TensorFlow Lite Implementation

We will use a pre-trained model (MobileNetV2) and convert it to TensorFlow Lite.

Loading and Saving the Model

The first thing that we do is import TensorFlow and load a pre-trained MobileNetV2 model. It is ready to utilize for pre-training on the ImageNet dataset, as has been seen in this model. The model.export (‘mobilenet_model’) writes the model in a format of TensorFlow’s SavedModel. This is the format required to convert it to the TensorFlow Lite Model (TFLite) that is used with mobile devices.

# Step 1: Set up the environment and load a pre-trained MobileNetV2 model
import tensorflow as tf

# Load a pretrained MobileNetV2 model
model = tf.keras.applications.MobileNetV2(weights='imagenet', input_shape=(224, 224, 3))

# Save the model as a SavedModel for TFLite conversion
model.export('mobilenet_model')

Convert the Model to TensorFlow Lite

The model is loaded from the saved model (mobilenet_model directory) using TFLiteConverter. The converter converts the model to a more lightweight .tflite format. Finally, the TFLite model is saved as mobilenet_v2.tflite for later use in mobile or edge applications.

# Step 2: Convert the model to TensorFlow Lite
converter = tf.lite.TFLiteConverter.from_saved_model('mobilenet_model')
tflite_model = converter.convert()

# Save the converted model to a TFLite file
with open('mobilenet_v2.tflite', 'wb') as f:
    f.write(tflite_model)

Loading the TFLite Model for Inference

Now, we import the necessary libraries for numerical operations (numpy) and image manipulation (PIL.Image). The TFLite model is loaded using tf.lite.Interpreter and memory are allocated for input/output tensors. We retrieve details about the input/output tensors, like the shapes and data types, which will be useful when we preprocess the input image and retrieve the output.

import numpy as np
from PIL import Image

# Load the TFLite model and allocate tensors
interpreter = tf.lite.Interpreter(model_path='mobilenet_v2.tflite')
interpreter.allocate_tensors()

# Get input and output tensors
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

Preprocessing Input, Running Inference, and Decoding Output

We load the image (cat.jpg), resize it to the required (224, 224) pixels, and preprocess it using MobileNetV2’s preprocessing method. The preprocessed image is fed into the TFLite model by setting the input tensor using interpreter.set_tensor(), and we run inference using interpreter.invoke(). After inference, we retrieve the model’s predictions and decode them into human-readable class names and probabilities using decode_predictions(). Finally, we print the predictions.

# Load and preprocess the input image
image = Image.open('cat.jpg').resize((224, 224))  # Replace with your image path
input_data = np.expand_dims(np.array(image), axis=0)
input_data = tf.keras.applications.mobilenet_v2.preprocess_input(input_data)

# Set the input tensor and run the model
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()

# Get the output and decode predictions
output_data = interpreter.get_tensor(output_details[0]['index'])
predictions = tf.keras.applications.mobilenet_v2.decode_predictions(output_data)
print(predictions)

Use the cat image below:

Output:

[ (‘n02123045’, ‘tabby’, 0.85), (‘n02124075’, ‘Egyptian_cat’, 0.07), (‘n02123159’, ‘tiger_cat’, 0.05)]

This means the model is 85% confident that the image is a tabby cat.

PyTorch Mobile Implementation

Now, we are going to implement PyTorch Mobile. We will use a simple pre-trained model like ResNet18, convert it to TorchScript, and run inference

Setting up the environment and loading the ResNet18 Model

# Step 1: Set up the environment
import torch
import torchvision.models as models

# Load a pretrained ResNet18 model
model = models.resnet18(pretrained=True)

# Set the model to evaluation mode
model.eval()

Converting the Model to TorchScript

Here, we define an example_input, which is a random tensor of size [1, 3, 224, 224]. This simulates a batch of 1 image with 3 color channels (RGB), and 224×224 pixels. It’s used to trace the model’s operations. torch.jit.trace() is a method that converts the PyTorch model into a TorchScript module. TorchScript allows you to serialize and run the model outside of Python, such as in C++ or mobile devices. The converted TorchScript model is saved as “resnet18_scripted.pt”, allowing it to be loaded and used later.

# Step 2: Convert to TorchScript
example_input = torch.randn(1, 3, 224, 224)  # Example input for tracing
traced_script_module = torch.jit.trace(model, example_input)

# Save the TorchScript model
traced_script_module.save("resnet18_scripted.pt")

Load the Scripted Model and Make Predictions

We use torch.jit.load() to load the previously saved TorchScript model from the file “resnet18_scripted.pt”. We create a new random tensor input_data, again simulating an image input with size [1, 3, 224, 224]. The model is then run on this input using loaded_model(input_data). This returns the output, which contains the raw scores (logits) for each class. To get the predicted class, we use torch.max(output, 1) which gives the index of the class with the highest score. We print the predicted class using predicted.item().

# Step 3: Load and run the scripted model
loaded_model = torch.jit.load("resnet18_scripted.pt")

# Simulate input data (a random image tensor)
input_data = torch.randn(1, 3, 224, 224)

# Run the model and get predictions
output = loaded_model(input_data)
_, predicted = torch.max(output, 1)
print(f'Predicted Class: {predicted.item()}')

Output:

Predicted Class: 107

Thus, the model predicts that the input data belongs to class index 107.

Conclusion

TensorFlow Lite gives more focus on mobile devices while PyTorch Mobile provides a more general CPU/GPU-deployed solution, both being optimized for the different applications of AI on mobile and edge devices. Compared to TensorFlow Lite, PyTorch Mobile offers greater portability while also being lighter than TensorFlow Lite and closely integrated with Google. Combined, they enable developers to implement real-time Artificial intelligence applications with high functionality on the developers’ handheld devices. These frameworks are empowering users with the capability to run sophisticated models on local machines and by doing so they are rewriting the rules for how mobile applications engage with the world, through fingertips.

Key Takeaways

• TensorFlow Lite and PyTorch Mobile empower developers to deploy AI models on edge devices efficiently.
• Both frameworks support cross-platform compatibility, enhancing the reach of mobile AI applications.
• TensorFlow Lite is known for performance optimization, while PyTorch Mobile excels in flexibility.
• Ease of integration and developer-friendly tools make both frameworks suitable for a wide range of AI use cases.
• Real-world applications span industries such as healthcare, retail, and entertainment, showcasing their versatility.

Frequently Asked Questions

The media shown in this article is not owned by Analytics Vidhya and is used at the Author’s discretion.