How to Use TensorFlow Lite for Mobile

TensorFlow Lite is Google's solution for running machine learning models on mobile and embedded devices. Unlike full TensorFlow, which prioritizes flexibility and training capabilities, TensorFlow...

Key Insights

  • TensorFlow Lite reduces model size by up to 75% through quantization while maintaining 95%+ accuracy, making it essential for mobile deployment where app size and memory constraints matter
  • GPU delegation can accelerate inference by 2-5x on mobile devices, but requires careful configuration and isn’t universally supported across all operations
  • Converting models to TensorFlow Lite format requires deliberate optimization decisions—post-training quantization is the easiest win, but for production apps you should benchmark dynamic range, float16, and int8 quantization variants

Introduction to TensorFlow Lite

TensorFlow Lite is Google’s solution for running machine learning models on mobile and embedded devices. Unlike full TensorFlow, which prioritizes flexibility and training capabilities, TensorFlow Lite is laser-focused on inference performance with minimal resource consumption.

The core value proposition is straightforward: take your trained TensorFlow or Keras model and convert it to a format that’s 4x smaller and runs significantly faster on mobile hardware. This matters because mobile apps operate under strict constraints—users won’t tolerate 100MB downloads or battery-draining computation.

TensorFlow Lite achieves this through aggressive optimization: operator fusion, quantization, and a stripped-down runtime that only includes what’s necessary for inference. The framework supports Android, iOS, Linux-based embedded systems, and microcontrollers. For mobile developers, Android and iOS are the primary targets.

The trade-off is reduced operator support compared to full TensorFlow. Not every TensorFlow operation has a TensorFlow Lite equivalent, though the coverage is comprehensive for common architectures like CNNs, RNNs, and transformers.

Converting Models to TensorFlow Lite Format

Model conversion is where you make critical optimization decisions. The TFLiteConverter handles the heavy lifting, but you need to understand quantization options.

Here’s a practical example converting a Keras image classification model with post-training quantization:

import tensorflow as tf
import numpy as np

# Load your trained Keras model
model = tf.keras.models.load_model('image_classifier.h5')

# Create converter
converter = tf.lite.TFLiteConverter.from_keras_model(model)

# Enable post-training quantization (dynamic range)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# For full integer quantization, provide representative dataset
def representative_dataset():
    # Load sample images from your training data
    for i in range(100):
        image = load_and_preprocess_image(f'sample_{i}.jpg')
        yield [np.array([image], dtype=np.float32)]

# Uncomment for int8 quantization
# converter.representative_dataset = representative_dataset
# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# converter.inference_input_type = tf.uint8
# converter.inference_output_type = tf.uint8

# Convert model
tflite_model = converter.convert()

# Save the model
with open('image_classifier.tflite', 'wb') as f:
    f.write(tflite_model)

# Verify conversion
interpreter = tf.lite.Interpreter(model_content=tflite_model)
print(f"Input details: {interpreter.get_input_details()}")
print(f"Output details: {interpreter.get_output_details()}")

Post-training quantization comes in three flavors: dynamic range (easiest, quantizes weights only), float16 (good balance), and int8 (maximum compression, requires representative data). For production, benchmark all three against your test set.

Setting Up TensorFlow Lite in Android

Android integration starts with Gradle dependencies. Add these to your app-level build.gradle:

dependencies {
    // TensorFlow Lite
    implementation 'org.tensorflow:tensorflow-lite:2.14.0'
    
    // GPU acceleration support (optional but recommended)
    implementation 'org.tensorflow:tensorflow-lite-gpu:2.14.0'
    
    // Support library for common preprocessing tasks
    implementation 'org.tensorflow:tensorflow-lite-support:0.4.4'
}

android {
    aaptOptions {
        noCompress "tflite"  // Prevent compression of model files
    }
}

Place your .tflite model file in app/src/main/assets/. Android Studio will package it with your APK.

For better organization, create a dedicated model manager class that handles initialization and resource cleanup.

Implementing Inference on Android

Here’s a complete Kotlin implementation for image classification:

import android.content.Context
import android.graphics.Bitmap
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.gpu.GpuDelegate
import java.io.FileInputStream
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.nio.MappedByteBuffer
import java.nio.channels.FileChannel

class ImageClassifier(context: Context) {
    private var interpreter: Interpreter
    private val labels: List<String>
    
    companion object {
        private const val MODEL_FILE = "image_classifier.tflite"
        private const val LABELS_FILE = "labels.txt"
        private const val INPUT_SIZE = 224
        private const val PIXEL_SIZE = 3
        private const val IMAGE_MEAN = 127.5f
        private const val IMAGE_STD = 127.5f
    }
    
    init {
        val model = loadModelFile(context)
        val options = Interpreter.Options().apply {
            setNumThreads(4)
            // GPU delegate for acceleration
            addDelegate(GpuDelegate())
        }
        interpreter = Interpreter(model, options)
        labels = loadLabels(context)
    }
    
    private fun loadModelFile(context: Context): MappedByteBuffer {
        val fileDescriptor = context.assets.openFd(MODEL_FILE)
        val inputStream = FileInputStream(fileDescriptor.fileDescriptor)
        val fileChannel = inputStream.channel
        return fileChannel.map(
            FileChannel.MapMode.READ_ONLY,
            fileDescriptor.startOffset,
            fileDescriptor.declaredLength
        )
    }
    
    private fun loadLabels(context: Context): List<String> {
        return context.assets.open(LABELS_FILE).bufferedReader().readLines()
    }
    
    fun classify(bitmap: Bitmap): List<Pair<String, Float>> {
        // Preprocess image
        val inputBuffer = preprocessImage(bitmap)
        
        // Prepare output buffer
        val outputBuffer = Array(1) { FloatArray(labels.size) }
        
        // Run inference
        interpreter.run(inputBuffer, outputBuffer)
        
        // Process results
        return outputBuffer[0]
            .mapIndexed { index, confidence -> labels[index] to confidence }
            .sortedByDescending { it.second }
            .take(5)
    }
    
    private fun preprocessImage(bitmap: Bitmap): ByteBuffer {
        val scaledBitmap = Bitmap.createScaledBitmap(bitmap, INPUT_SIZE, INPUT_SIZE, true)
        val buffer = ByteBuffer.allocateDirect(4 * INPUT_SIZE * INPUT_SIZE * PIXEL_SIZE)
        buffer.order(ByteOrder.nativeOrder())
        
        val pixels = IntArray(INPUT_SIZE * INPUT_SIZE)
        scaledBitmap.getPixels(pixels, 0, INPUT_SIZE, 0, 0, INPUT_SIZE, INPUT_SIZE)
        
        for (pixel in pixels) {
            val r = ((pixel shr 16 and 0xFF) - IMAGE_MEAN) / IMAGE_STD
            val g = ((pixel shr 8 and 0xFF) - IMAGE_MEAN) / IMAGE_STD
            val b = ((pixel and 0xFF) - IMAGE_MEAN) / IMAGE_STD
            
            buffer.putFloat(r)
            buffer.putFloat(g)
            buffer.putFloat(b)
        }
        
        return buffer
    }
    
    fun close() {
        interpreter.close()
    }
}

This implementation handles the complete pipeline: model loading, image preprocessing with normalization, inference execution, and result parsing. The GPU delegate is enabled by default for hardware acceleration.

iOS Integration with TensorFlow Lite

For iOS, use CocoaPods or Swift Package Manager. Here’s the Podfile approach:

pod 'TensorFlowLiteSwift'
pod 'TensorFlowLiteSelectTfOps'  # If you need select TF ops

Swift implementation follows similar patterns:

import TensorFlowLite
import UIKit

class ImageClassifier {
    private var interpreter: Interpreter
    private let labels: [String]
    
    private let inputWidth = 224
    private let inputHeight = 224
    private let batchSize = 1
    private let inputChannels = 3
    
    init?(modelFileName: String = "image_classifier") {
        guard let modelPath = Bundle.main.path(
            forResource: modelFileName,
            ofType: "tflite"
        ) else {
            print("Failed to load model")
            return nil
        }
        
        do {
            var options = Interpreter.Options()
            options.threadCount = 4
            
            interpreter = try Interpreter(modelPath: modelPath, options: options)
            try interpreter.allocateTensors()
            
            // Load labels
            guard let labelsPath = Bundle.main.path(forResource: "labels", ofType: "txt"),
                  let labelsContent = try? String(contentsOfFile: labelsPath) else {
                return nil
            }
            labels = labelsContent.components(separatedBy: .newlines)
        } catch {
            print("Interpreter initialization failed: \(error)")
            return nil
        }
    }
    
    func classify(image: UIImage) -> [(String, Float)]? {
        guard let pixelBuffer = image.pixelBuffer(width: inputWidth, height: inputHeight) else {
            return nil
        }
        
        let inputData = preprocessImage(pixelBuffer)
        
        do {
            try interpreter.copy(inputData, toInputAt: 0)
            try interpreter.invoke()
            
            let outputTensor = try interpreter.output(at: 0)
            let results = [Float](unsafeData: outputTensor.data) ?? []
            
            return zip(labels, results)
                .sorted { $0.1 > $1.1 }
                .prefix(5)
                .map { ($0.0, $0.1) }
        } catch {
            print("Inference failed: \(error)")
            return nil
        }
    }
    
    private func preprocessImage(_ pixelBuffer: CVPixelBuffer) -> Data {
        // Normalize pixel values to [-1, 1]
        var data = Data()
        // Implementation details omitted for brevity
        return data
    }
}

Performance Optimization and Best Practices

GPU delegation provides the biggest performance boost, but configuration matters:

// Android GPU configuration
val options = Interpreter.Options()
val gpuDelegate = GpuDelegate(GpuDelegate.Options().apply {
    setPrecisionLossAllowed(true)  // Allow FP16 for better performance
    setInferencePreference(GpuDelegate.Options.INFERENCE_PREFERENCE_SUSTAINED_SPEED)
})
options.addDelegate(gpuDelegate)

// Enable NNAPI for Android (alternative to GPU)
val nnApiDelegate = NnApiDelegate()
options.addDelegate(nnApiDelegate)

// Thread configuration for CPU fallback
options.setNumThreads(Runtime.getRuntime().availableProcessors())

Key optimization strategies:

  1. Use GPU delegation for models with many convolutional operations
  2. Enable quantization during conversion—int8 models run faster and use less memory
  3. Batch preprocessing if running multiple inferences
  4. Cache the interpreter—initialization is expensive, reuse instances
  5. Profile on real devices—emulators don’t reflect actual performance

Testing and Deployment Considerations

Always benchmark your model on target hardware:

class InferenceBenchmark(private val classifier: ImageClassifier) {
    fun benchmark(testImages: List<Bitmap>, iterations: Int = 100): BenchmarkResults {
        val latencies = mutableListOf<Long>()
        
        // Warmup
        repeat(10) { classifier.classify(testImages.first()) }
        
        // Actual benchmark
        testImages.forEach { image ->
            repeat(iterations) {
                val start = System.nanoTime()
                classifier.classify(image)
                val end = System.nanoTime()
                latencies.add((end - start) / 1_000_000) // Convert to ms
            }
        }
        
        return BenchmarkResults(
            mean = latencies.average(),
            p50 = latencies.sorted()[latencies.size / 2],
            p95 = latencies.sorted()[(latencies.size * 0.95).toInt()],
            p99 = latencies.sorted()[(latencies.size * 0.99).toInt()]
        )
    }
}

Before deploying, verify model accuracy on-device matches your validation results. Quantization can introduce small accuracy degradations—measure them. Test edge cases: poor lighting, unusual angles, low-resolution inputs.

For app size optimization, consider model compression techniques beyond quantization: weight sharing, pruning, or knowledge distillation if you control the training pipeline.

TensorFlow Lite enables powerful on-device ML, but success requires careful optimization and testing. Start with post-training quantization, enable GPU delegation, and always benchmark on real hardware before shipping.

Liked this? There's more.

Every week: one practical technique, explained simply, with code you can use immediately.