MTK Converter

Overview

This page provides a comprehensive guide for converting YOLOv5s PyTorch models to TensorFlow Lite (TFLite) format using the MTK Converter tool. It covers the setup of the YOLOv5 environment, the necessary patches, and detailed steps for converting models to both int8 and fp32 TFLite formats.

YOLOv5s Model Conversion

• Source Link: http://github.com/ultralytics/yolov5

• Model Name: YOLOv5s

Setting Up YOLOv5 Environment

Before you begin, ensure that the NeuroPilot Converter Tool is installed.

1. Clone the repository:

   git clone http://github.com/ultralytics/yolov5
   cd yolov5
   git reset --hard 485da42
   

2. Install Python packages and dependencies:

   pip3 install -r requirements.txt
   pip3 install torch==1.9.0 torchvision==0.10.0
   

   Note

   The mtk_converter.PyTorchConverter only supports PyTorch versions between 1.3.0 and 2.0.0. The detected version v2.3.1+cu121 is not within this supported range, causing a runtime error. Therefore, it is necessary to install a compatible version of PyTorch and torchvision to ensure compatibility.

3. Apply Patch:

   git apply Fix_yolov5_mtk_tflite_issue.patch
   

   Note

   This patch adds support for MTK TensorFlow Lite (MTK TFLite) in the YOLOv5 model export script. It includes:

   • Adding mtk_tflite as a supported export format.

   • Modifying the Detect module’s forward method to only include convolution operations.

   • Implementing post-processing operations for MTK TFLite.

   • Extending the DetectMultiBackend class to handle MTK TFLite models.

4. Exporting Pytorch Model to TorchScript Format:

   python export.py --weight yolov5s.pt --img-size 640 640 --include torchscript
   

YOLOv5 PyTorch Model Conversion to int8/fp32 TFLite Formats

• Quant8 Format:

1. Prepare Calibration Data:

To prepare the calibration data, create a new Python script named prepare_calibration_data.py in the root directory of YOLOv5 project. This script will generate a set of images that are used for model quantization calibration.

python prepare_calibration_data.py

import os
import numpy as np
from utils.dataloaders import LoadImagesAndLabels
from utils.general import check_dataset

data = 'data/coco128.yaml'
num_batches = 100
calib_dir = 'calibration_dataset'
os.makedirs(calib_dir)

# Retrieve first 100 images from training set with batch_size = 1
dataset = LoadImagesAndLabels(check_dataset(data)['train'], batch_size=1)

for idx, (im, _target, _path, _shape) in enumerate(dataset):
    if idx >= num_batches:
        break

    # Expand shape from (3, 640, 640) to (1, 3, 640, 640)
    im = np.expand_dims(im, axis=0).astype(np.float32)
    # 0 - 255 to 0.0 - 1.0
    im /= 255
    np.save(os.path.join(calib_dir, 'batch-{:05d}.npy'.format(idx)), im)

2. Convert to int8 TFLite:

To perform the conversion of the Pytorch model to an int8 TFLite format, create a new Python script named convert_to_quant_tflite.py in the root directory of your YOLOv5 project. This script will handle the conversion process by utilizing the pre-generated calibration data and converting the model into the quantized TFlite format.

python convert_to_quant_tflite.py

import os
import numpy as np
import mtk_converter

calib_dir = 'calibration_dataset'

converter = mtk_converter.PyTorchConverter.from_script_module_file(
    'yolov5s.torchscript', input_shapes=[(1, 3, 640, 640)]
)

def data_gen():
    """Return an iterator for the calibration dataset."""
    for fn in sorted(os.listdir(calib_dir)):
        yield [np.load(os.path.join(calib_dir, fn))]

converter.quantize = True
converter.calibration_data_gen = data_gen
converter.convert_to_tflite('yolov5s_int8_mtk.tflite')

• FP32 Format:

1. Convert to fp32 TFLite:

To convert the Pytorch model to an fp32 TFLite format, create a new Python script named convert_to_tflite.py in the root directory of your YOLOv5 project. This script will handle the conversion process to generate a non-quantized, full-precision TFLite model.

python convert_to_tflite.py

import mtk_converter

converter = mtk_converter.PyTorchConverter.from_script_module_file(
    'yolov5s.torchscript', input_shapes=[(1, 3, 640, 640)]
)
converter.convert_to_tflite('yolov5s_mtk.tflite')