Train a TFlite model

Many machines have cameras through which they can monitor their environment. With machine leaning (ML), you can train models on patterns within image data. Follow this guide to use your image data to train an ML model, so that your machine can make inferences about its environment.

Prerequisites

a machine connected to the Viam app
Add a new machine in the Viam app. On the machine’s page, follow the setup instructions to install viam-server on the computer you’re using for your project. Wait until your machine has successfully connected to the Viam app.
a dataset with labels
Follow the guide to create a dataset.

Train a machine learning model

Now that you have a dataset that contains your labeled images, you are ready to train a machine learning model.

1

Find your training dataset

In the Viam app, navigate to your list of DATASETS and select the one you want to train on.

2

Train an ML model

Click Train model and follow the prompts. You can train a TFLite model using Built-in training.

The shapes dataset.

Click Next steps.

3

Fill in the details for your ML model

Enter a name for your new model.

Select a Task Type:

  • Single Label Classification: The resulting model predicts one of the selected labels or UNKNOWN per image. Select this if you only have one label on each image. Ensure that the dataset you are training on also contains unlabeled images.
  • Multi Label Classification: The resulting model predicts one or more of the selected labels per image.
  • Object Detection: The resulting model predicts either no detected objects or any number of object labels alongside their locations per image.

Select the labels you want to train your model on from the Labels section. Unselected labels will be ignored, and will not be part of the resulting model.

Click Train model.

The data tab showing the train a model pane
4

Wait for your model to train

The model now starts training and you can follow its process on the TRAINING tab.

Once the model has finished training, it becomes visible on the MODELS tab.

You will receive an email when your model finishes training.

5

Debug your training job

From the TRAINING tab, click on your training job’s ID to see its logs.

You can also view your training jobs’ logs with the viam train logs command.

Test your ML model

Once your model has finished training, you can test it.

Ideally, you want your ML model to be able to work with a high level of confidence. As you test it, if you notice faulty predictions or confidence scores, you will need to adjust your dataset and retrain your model.

If you trained a classification model, you can test it with the following instructions.

  1. Navigate to the DATA page and click on the Images subtab.
  2. Click on an image to open the side menu, and select the Actions tab.
  3. In the Run model section, select your model and specify a confidence threshold.
  4. Click Run model.

If the results exceed the confidence threshold, the Run model section shows a label and the responding confidence threshold.

You can test both detection models and classifier models using the following resources together:

Iterate on your ML model

With the right training dataset, your ML model can perform well as soon as training completes. However, edge cases always exist. You are unlikely to account for all false positives or false negatives during your first round of training. Instead, keep track of edge cases as you run your model. Add images capturing edge cases to your training dataset, annotate them, and re-train your model using the new data. Using this approach, each subsequent model version becomes more accurate than the last.

To capture images of edge cases and re-train your model using those images, complete the following steps:

  1. Add edge case images to your training dataset. You can find edge cases in your existing data on the DATA page or capture new images and add them to your training dataset.

  2. Visit the DATASET tab of the DATA page and annotate the image.

  3. Repeat the steps above to train and release a new version of your ML model. Your machines will automatically update to the new version of the model soon after release.

Next steps

Now your machine can make inferences about its environment. The next step is to deploy the ML model and then act or alert based on these inferences.

See the following tutorials for examples of using machine learning models to make your machine do things based on its inferences about its environment: