Running Tflite Model Maker in Google Colab June/July 2023

I thought I’d quickly share/run through how I’m currently running Tflite Model Maker in Google Colab in June and July 2023 after Colab has updated to Python 3.10 and Tflite Model Maker isn’t currently compatible due to TF Support Addons dependency.

This is all based off the Github comment thread on the issue and is definitely a hacky way around the issue, but for all those wanting a temporary solution here goes.

The basic deal is you have to install a Python 3.9 environment on Google Colab and force it to run all your code through that environment. I’ll base this demo off the Android Figurine tutorial but you can just amend to your own needs/incorporate your own dataset.

There are 3 major parts to the pipeline:

  1. Environment Install
  2. Data Uploading
  3. Train and Evaluation Script Running

Before you get going you need to build a Training and Evaluation Python Script that you will upload and run in Colab. You’ll do this towards the end it’s just to have it handy now.

This python scipt is basically your old Tflite Model Maker cells without the installs.

You can build the below script in Colab then click File > Download as .py file (if you want to customise it/use your old training pipeline)

Save this file as train.py (you’ll upload it to Colab later in the process)

# -*- coding: utf-8 -*-
""" train.py

Automatically generated by Colaboratory.

"""

# Imports
import numpy as np
import os

from tflite_model_maker.config import ExportFormat, QuantizationConfig
from tflite_model_maker import model_spec
from tflite_model_maker import object_detector

from tflite_support import metadata

import tensorflow as tf
assert tf.__version__.startswith('2')

tf.get_logger().setLevel('ERROR')
from absl import logging
logging.set_verbosity(logging.ERROR)

# Confirm TF Version
print("\nTensorflow Version:")
print(tf.__version__)
print()


# Load Dataset
train_data = object_detector.DataLoader.from_pascal_voc(
    'android_figurine/train',
    'android_figurine/train',
    ['android', 'pig_android']
)

val_data = object_detector.DataLoader.from_pascal_voc(
    'android_figurine/validate',
    'android_figurine/validate',
    ['android', 'pig_android']
)

# Load model spec
spec = object_detector.EfficientDetSpec(
  model_name='efficientdet-lite2',
  uri='https://tfhub.dev/tensorflow/efficientdet/lite2/feature-vector/1',
  model_dir='/content/checkpoints',
  hparams={'max_instances_per_image': 8000})

# Train the model
model = object_detector.create(train_data, model_spec=spec, batch_size=4, train_whole_model=True, epochs=20, validation_data=val_data)

# Evaluate the model
eval_result = model.evaluate(val_data)

# Print COCO metrics
print("COCO metrics:")
for label, metric_value in eval_result.items():
    print(f"{label}: {metric_value}")

# Add a line break after all the items have been printed
print()

# Export the model
model.export(export_dir='.', tflite_filename='android.tflite')

# Evaluate the tflite model
tflite_eval_result = model.evaluate_tflite('android.tflite', val_data)

# Print COCO metrics for tflite
print("COCO metrics tflite")
for label, metric_value in tflite_eval_result.items():
    print(f"{label}: {metric_value}")

The Colab Workbook

There are a few steps that require you to type Y to proceed with installation. So run one cell at a time and monitor until training begins.

  1. Open New Workbook in Colab

  2. Change Runtime to GPU

  3. Create the following cells, then once ALL are created, run them one at a time.

%env PYTHONPATH = # /env/python
# Download and install and update miniconda (you will need to hit y to proceed) #
!wget https://repo.anaconda.com/miniconda/Miniconda3-py39_23.3.1-0-Linux-x86_64.sh
!chmod +x Miniconda3-py39_23.3.1-0-Linux-x86_64.sh
!./Miniconda3-py39_23.3.1-0-Linux-x86_64.sh -b -f -p /usr/local
!conda update -q conda
import sys
sys.path.append('/usr/local/lib/python3.9/site-packages')
# create conda environment (you will need to hit y to proceed) #
!conda create -n myenv python=3.9
# Download and unzip your data files organised appropriately #
!wget https://storage.googleapis.com/download.tensorflow.org/data/android_figurine.zip
!unzip -q android_figurine.zip
# Upload your own train.py (to /content/) or the one we created above or download this existing one
!wget https://raw.githubusercontent.com/wwfish/tflite-model-maker/main/train.py
# Install dependencies in the new environment. You need to force activate conda every time
%%shell
eval "$(conda shell.bash hook)"
conda activate myenv
pip install -q tflite-model-maker
pip3 install -q pycocotools
pip install -q ipykernel
pip install -q numpy==1.23.4
# Run your training and evaluation script as an external python script
%%shell
eval "$(conda shell.bash hook)"
conda activate myenv
python /content/train.py

There are lots of moving parts going on here and Google Colab keeps returning back to their default environment, thats why you have to keep activating your conda environment over and over, hence why it’s all a bit sketchy. But as of tonight this was working for me.

Tflite file will be created in /content/ directory of Colab.

Have updated train.py code to print the COCO metrics for eval_result and tflite_eval_result

This all comes courtesy of tomkuzma over on Github.

Happy for people to add any updates/changes/better things to this.

3 Likes

Put the Colab workbook and the train.py file in a Git for people who want to make things a bit easier.

Tflite-Model-Maker-Workaround Git.

2 Likes

Thank you so much @wwfisher for this workaround!