Process for training Tflite Model Maker (EfficientDet) in Google Colab in June/July 2023 with tflite-model-maker not currently being compatible with current version of Colab. This example is based upon the Android Figurine Colab workbook published here.
This code is based off the work of @tomkuzma and his work in the comment thread here.
Google Colab has updated to Python 3.10 which has made the Tflite Model Maker API stop working/load on an infinite loop.
This Colab workbook loads a Conda environment into Google Colab with Python 3.9 and allows for the Tflite Model Maker to be installed and trained.
The working demo is based upon the Android Figurine training example from here: https://colab.research.google.com/github/khanhlvg/tflite_raspberry_pi/blob/main/object_detection/Train_custom_model_tutorial.ipynb
Edit train.py script as needed for your own training requirements/dataset.
The process of how to do this is outlined in this post:
https://discuss.tensorflow.org/t/running-tflite-model-maker-in-google-colab-june-july-2023/18049
- Edit/Build a train.py config file to suit your dataset and your training specifications - download it to use in your own Colab workbook.
- Run the Colab workbook to install tflite-model-maker and run the training file (remember to upload your own train.py file to Colab).
Copy and paste the following code block into a .py file ready to upload the Colab workbook and run.
In this train.py you can edit:
- Model architecture specification
- Folder structure for train, valid and test data
- Batch size for training
- Epoch amounts for training
# -*- 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}")
In this Colab workbook you will get:
- a .tflite file trained
- COCO metrics for unquantized model
- COCO metrics for quantized model
- Training checkpoints data in checkpoints folder
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.
Because of the way Google Colab works you need to activate your conda environment with the %%shell command everytime you want to run any command in the environment you've installed (Python 3.9). If you don't it will revert back to the default Colab environment (Python 3.10).
-
Open New Workbook in Colab
-
Change Runtime to GPU
-
Create the below code blocks as their own cells, then once ALL are created, run them one at a time.
-
Remember to upload your own train.py (code block 5) to meet your own requirements.
-
If building your own dataset remember to get the folder and file structure correct before uploading to Colab and training (code block 4). This example uses Pascal VOC xmls in the dataloader.
%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
Tensorflow team has pointed out that they are making updates to tflite-model-maker to work with Python 3.10 before they shift over to using Mediapipe as the go to training environment for tflite models.