Building a YOLOv3 pipeline with Valohai and Superb AI

Superb AI Inc. company logo

Superb AI

2021/6/14 | 4 min read

This article shows an example of a pipeline that integrates Valohai and Superb AI to train a computer vision model using pre-trained weights and transfer learning. For the model, we are using YOLOv3, which is built for real-time object detection.

Superb AI is used in this example to manage the training data and produce a labeled dataset that can be used for training. Valohai, on the other hand, takes in that data and runs a pipeline with a training step and a batch inference step. We’ve written a more thorough breakdown of how these platforms complement each other here.

  • Superb AI is an ML DataOps platform that helps computer vision teams automate and manage the full data pipeline: from ingestion and labeling to data quality assessment and delivery.

  • Valohai is an MLOps platform that helps build and operate fully automated training pipelines on any cloud or on-prem hardware. Perfect for computer vision and deep learning.

Project details

Note: The pipeline is generic, which means it is trivial to train with any dataset and any kind of classes from the Superb.ai platform, using any weights that fit the YOLOv3 model. Listed here are just the defaults we used.

Preparing the data in Superb AI

First, we need to create a new project in the Superb AI platform. Choose any name and description. The data type will be Image and annotation type Box for this example.

Then add five classes (person, car, chair, book, bottle) to match our COCO dataset, which we will upload to the platform.

One can utilize the Superb AI’s automatic labeling features too, but in this example, we will have the labels already in the COCO dataset. Now that we have our project and dataset template prepared in the platform, it is time to convert and upload the actual data + labels using Superb AI SDK and CLI. Follow the instructions in this example repository.

Once the conversion and upload are done, the project should contain more than 3000 images, ready to be used by the training pipeline in Valohai.

The Valohai training pipeline

In Valohai, the training pipeline is a key component of any successful project. Our example pipeline consists of four steps.

The convert-superbai step pulls the data from Superb.ai, splits the training and validation sets based on Superb.ai tagging, and finally converts the two sets into .tfrecord files for the model.

The weights step downloads the pre-trained weights from Joseph Redmon’s website, converts them to a Tensorflow checkpoint, and finally validates that they work correctly by running inference for a single image.

The train step trains the model in a GPU environment (other steps are CPU by default) based on the hyperparameters.

Finally, the detect step runs inference on the freshly trained model for the set images given as an input for the step. The output of this step is new versions of the images, with bounding boxes and labels rendered on top of the original image.

Each step is (usually) a single Python file. Here is the source code for weights.py of our weights step. It has one parameter weights_num_classes (how many classes were used for the pre-trained model) and one input weights (the actual pre-trained weights).

import numpy as np

from yolov3_tf2.models import YoloV3
from yolov3_tf2.utils import load_darknet_weights
import tensorflow as tf
import valohai

params = {
   "weights_num_classes": 80,
}

inputs = {
    "weights": "https://pjreddie.com/media/files/yolov3.weights",
}

valohai.prepare(step="weights", default_parameters=params, default_inputs=inputs)

physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
   tf.config.experimental.set_memory_growth(physical_devices[0], True)

yolo = YoloV3(classes=valohai.parameters('weights_num_classes').value)
load_darknet_weights(yolo, valohai.inputs('weights').path(), False)
# Sanity check with random image
img = np.random.random((1, 320, 320, 3)).astype(np.float32)
output = yolo(img)
path = valohai.outputs('model').path('model.tf')
yolo.save_weights(path)

The parameters and inputs are fed to the valohai.prepare() call to Valohai SDK. Valohai CLI can parse this call to automatically generate the necessary configuration for parameterizing the step.

The pipeline itself is defined in the pipeline.py, which also uses the new SDK.

import valohai

def main(old_config):
   pipeline = valohai.Pipeline(name="train-superbai", config=old_config)

# Define nodes
   convert = papi.execution("convert-superbai")
   weights = papi.execution("weights")
   train = papi.execution("train")
   detect = papi.execution("detect")

# Configure pipeline
   convert.output("classes.txt").to(train.input("classes"))
   convert.output("train/*").to(train.input("train"))
   convert.output("test/*").to(train.input("test"))
   convert.output("classes.txt").to(detect.input("classes"))
   weights.output("model/*").to(train.input("model"))
   train.output("model/*").to(detect.input("model"))

return pipeline

Setting up the pipeline in Valohai

First, create a new project in Valohai and link the GitHub example repository from the project settings.

Now you can create a new pipeline from the pipelines tab and select train-superbai from the blueprints dropdown.

Integrating with Superb AI

Before running the pipeline, you need to tell it where to get the data. The images and labels are two separate URLs, which means two copy-pastes from Superb.ai to Valohai.

Images: Settings > General (Superb.ai)

Labels: Export > Export History (Superb.ai)

Pipeline Inputs (Valohai)

Running the pipeline

Now that we are integrated with our dataset in the Superb AI platform, our pipeline is ready to rumble!

Valohai records all the inputs, parameters, logs, and metrics for each step of the pipeline. These are the inputs, parameters, Docker image, environment, and the rest of the information for the execution of the train step:

This is a slice of all the logs printed out by the training code.

These are the metrics printed out by the training code. Valohai parses all printed JSON as a metric (called metadata in Valohai).

Finally, in the outputs tab of the detect step, we can inspect how the model is performing with some actual images.

Iterating on labels

Once you know where your model needs optimization as it relates to data, you can jumpstart your next iteration by quickly sending your model’s output back into the Superb AI Suite via their SDK for model-assisted labeling.

Why setting up an end-to-end pipeline for your YOLOv3 model is a no-brainer

In this article, the end-to-end pipeline is built on top of two purpose built tools: Superb AI for the data management and Valohai for the model lifecycle management. The whole setup in this example can be done within hours, but the value of the upfront work compounds every time the model is trained.

The training dataset evolves over time and it can be relabeled and changed independently from the training pipeline. Similarly, the code in the training pipeline can be iterated, while keeping the end-to-end pipeline intact for continuous operations. Additionally, the training pipeline on Valohai can be scaled with ease to more powerful GPU-instances as the training dataset and modeling complexity grows.

Subscribe to our newsletter

Stay updated latest MLOps news and our product releases

About Superb AI

Superb AI is an enterprise-level training data platform that is reinventing the way ML teams manage and deliver training data within organizations. Launched in 2018, the Superb AI Suite provides a unique blend of automation, collaboration and plug-and-play modularity, helping teams drastically reduce the time it takes to prepare high quality training datasets. If you want to experience the transformation, sign up for free today.

Join The Ground Truth Community

The Ground Truth is a community newsletter featuring computer vision news, research, learning resources, MLOps, best practices, events, podcasts, and much more. Read The Ground Truth now.

home_ground_truth

Designed for Data-Centric Teams

We’ve built a platform for everyone involved in the journey from training to production - from data scientists and engineers to ML engineers, product leaders, labelers, and everyone in between. Get started today for free and see just how much faster you can go from ideation to precision models.