Converting YOLOv5 PyTorch to TensorFlow Lite Format

Converting YOLOv5 PyTorch Model Weights to TensorFlow Lite Format

In this article we test a face mask detector on a regular computer.
Here we make our model understandable to TensorFlow Lite, the lightweight version of TensorFlow specially developed to run on small devices.


In the previous article of this series, we trained and tested our YOLOv5 model for face mask detection. In this one, we’ll convert our model to TensorFlow Lite format.

I previously mentioned that we’ll be using some scripts that are still not available in the official Ultralytics repo (clone this) to make our life easier. To perform the transformation, we’ll use the script, which simplifies the PyTorch to TFLite conversion. Otherwise, we’d need to stick to the Ultralytics-suggested method that involves converting PyTorch to ONNX to TensorFlow to TFLite. Note that the last operation can fail, which is really frustrating.

To perform the conversion, run this:

import tensorflow as tf
!PYTHONPATH=. python3 /content/yolov5/models/ --weight
/content/yolov5/runs/train/exp2/weights/ --cfg models/yolov5s.yaml --img 416 --source /content/FaceMaskDataset/images/train

When passing the weights’ file path (the configuration.yaml file), indicate the image dimensions the model accepts and the source of the training dataset (the last parameter is optional). The script will use TensorFlow 2.3.1 to transform the .pt weights to the TensorFlow format and the output will be saved at /content/yolov5/runs/train/exp/weights.

Ultralytics-Suggested Method for Weights Conversion

From my perspective, this step is a bit cumbersome, but it’s necessary to show how it works. Instead of running the previous commands, run these lines:

# Uncomment all this if you want to follow the long path
#!pip install onnx>=1.7.0 # for ONNX export
#!pip install coremltools==4.0 # for CoreML export
#!python models/ --weights /content/yolov5/runs/train/exp2/weights/ --img 416 --batch 1 # export at 640x640 with batch size 1
#!pip install tensorflow-addons
#!pip install onnx-tf
#!pip install tensorflow==2.3.0
#import tensorflow_addons as tfa
#import onnx
#from onnx_tf.backend import prepare
#import tensorflow as tf
#base_model = onnx.load('/content/yolov5/runs/train/exp2/weights/best.onnx')
#to_tf = prepare(base_model)
#converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model('/content/yolov5/runs/train/exp2/weights/customyolov5')
#tflite_model = converter.convert() #just FYI: this step could go wrong and your notebook instance could crash.

Testing TFLite Weights on Google Colab

Now it’s time to check if the weights conversion went well. Before doing so, we need to slightly modify the script and set the proper class names. Open up the file (/content/yolov5/, look for names = [...] on line 157 and change it to names = ['Face mask','No face mask']. Save and close the file.

If everything went well, you should be able to load and test what you've obtained. Run the lines below. They will load the YOLOv5 model with the .tflite weights and run detection on the images stored at /test_images.

!python --weight /content/yolov5/runs/train/exp2/weights/best-fp16.tflite --img 416 --source ../test_images

If all goes well, the result will be similar to this:

Image 1

And with that, you're done — at least in this Notebook! As a last step, download the weights file stored at /content/yolov5/runs/train/exp/weights/best-fp16.tflite and to use them in the real-world implementation.

TFLite Interpreter for Good Performance on Edge Devices

The YOLOv5’s script uses a regular TensorFlow library to interpret TensorFlow models, including the TFLite formatted ones. In our scenario, TensorFlow is too heavy and resource-demanding to be run on small devices. If you want to maintain good performance of detections, better stick to TFLite and its interpreter. I’ve essentially replaced all TensorFlow-related operations with their TFLite equivalents. In addition, I made some small changes to make the detector able to run on TPU/GPU: I copied the file, modified it, and saved it as You can find the file here. I invite you to compare these files to fully understand the modifications.

Testing the YOLOv5 Model Weights Locally

This step is optional but recommended. In this short test, I’ll show you how to feed your computer’s webcam output to the detector before the final deployment on Pi. I’ll also show you how to test the model with and without the TFLite interpreter. Once you’ve got the modified file, create a folder on your local computer with the name Face Mask Detection.

Get the latest PyTorch version and its dependencies by running pip3 install torch torchvision from any CLI window.

Launch a Jupyter Notebook from the directory you’ve created: open the CLI, navigate to that folder, and issue the jupyter notebook command. Once the notebook pops up, run the following cells:

#Install the appropriate tensorflow version
!pip install tensorflow==2.3.1
import torch #It's required by YOLO
import tensorflow as tf
!git clone #comment this if this is not your first run
%cd yolov5
!git checkout tf-android
#Install all dependencies indicated at requirements.txt file
#Prompt if any GPU is available
!pip install -r requirements.txt
print('All set. Using PyTorch version %s with %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

Before continuing, remember to modify names list at line 157 in the file and copy all the downloaded weights into the /weights folder within the YOLOv5 folder.

To feed your YOLOv5 model with the computer’s webcam, run this command in a new notebook cell:

!python --weights weights/ --img 416 --conf 0.4 --source 0 --classes 0 1

It will initiate the webcam in a separate window, identify your face, and detect if you’re wearing a face mask or not. This is what you should expect:Image 2

Image 3

If you want to test the model with its TFLite weights, you first need to install the corresponding interpreter on your machine.

If you’re using linux x86_64:


If you’re using any other OS, I would suggest you check the best version for you.

Now you can run the next cell and expect exactly the same result as before:

!python --weights weights/best-fp16.tflite --img 416 --conf 0.45 --source 0 --classes 0 1

Next Step

We’ve trained and tested the YOLOv5 face mask detector. In the next article, we’ll deploy it on Raspberry Pi as promised. Stay tuned!

This article is part of the series 'AI on the Edge: Face Mask DetectionView All


This article, along with any associated source code and files, is licensed under The Code Project Open License (CPOL)

이 글을 페이스북으로 퍼가기 이 글을 트위터로 퍼가기 이 글을 카카오스토리로 퍼가기 이 글을 밴드로 퍼가기

Computer Vision

번호 제목 글쓴이 날짜 조회수
18 Computer Vision Annotation Tool (CVAT) 관리자 06-22 482
17 모델 파일 속성 확인 관리자 05-27 548
16 Converting YOLOv5 PyTorch to TensorFlow Lite Format 관리자 05-27 745
15 TensorFlow와 PyTorch 동시 활용하기 관리자 05-27 532
14 코랩 런타임 관리자 05-27 1,246
13 Resnet 관리자 05-27 978
12 대장 내시경 관리자 05-21 428
11 BERT를 활용하여 한국어 사전학습모델 KR-BERT 관리자 05-20 375
10 NER - 개체명 인식 관리자 05-19 856
9 Winodws / Linux 플랫폼에 다양한 버전의 pytorch를 Pip 설치 관리자 05-17 340
8 yolo 관리자 05-13 431
7 Deep Learning to Improve Breast Cancer Detection 관리자 05-11 447
6 opencv로 좌표 형태의 도형을 mask 이미지 배열로 변환, 또 역으로 관리자 05-09 399
5 Mask-RCNN 수행하기 - OpenCV DNN 모듈 관리자 05-09 456
4 Opencv a to z 관리자 03-14 417
3 Opencv Deep learning 이미지 인식 관리자 03-04 389
2 Opencv Deep learning 관리자 03-04 407
1 안드로이드 학습관련 관리자 10-26 772