Train a vision model with built-in training

Train an image classification or object detection model from a dataset of labeled images, with no training code. Viam runs the job on cloud GPUs and saves the trained model in your organization’s registry. To train any other kind of model, such as a robot policy or a PyTorch model, use a custom training script; to choose between the two, see Built-in or custom training.

Prerequisites

Built-in training rejects a job if the dataset doesn’t meet these minimums:

RequirementMinimum
Labeled images80% of the dataset
Examples per label10 images (classification) or 10 bounding boxes (object detection)
Total images15 (object detection only; classification has no total-image floor)

Built-in training uses only the images in a dataset and skips any other files, such as video or audio. It can’t train on a sequence dataset.

Choose a framework

The framework sets which software runs the model on your machine. Built-in training produces either one:

TFLiteTensorFlow
Choose it whenYou aren’t sure, or the model runs on a small computer such as a Raspberry PiYour machine has an NVIDIA GPU you want the model to use
Runs on the machine intflite_cpu, on the CPU of almost any machinetensorflow-cpu, or triton on an NVIDIA GPU
ResourcesMade for devices with little memory and computeNeeds more memory and compute, and takes longer to train

For the platforms each service supports, see What runs the model.

1. Start a training job

  1. In the Viam app, go to DATA > DATASETS and click the dataset you want to train on.
  2. Click Train model.
  3. Select New model. To retrain a model you already have, select New version and pick the model instead.
  4. Select Built-in (TensorFlow Lite) or Built-in (TensorFlow).
  5. Enter a Name for the model, such as part-inspector. The name identifies the model in your organization’s registry.
  6. Select the task type: Single label classification, Multi label classification, or Object detection. It must match how you labeled the dataset.
  7. Select the labels to train on. Multi label classification needs at least two. Leave out labels with too few examples or that you don’t want the model to learn.
  8. Click Train model.
viam train submit managed \
  --dataset-id=<dataset-id> \
  --model-org-id=<org-id> \
  --model-name=part-inspector \
  --model-type=single_label_classification \
  --model-framework=tflite \
  --model-labels=good-part,defective-part

Set --model-type to single_label_classification, multi_label_classification, or object_detection, and --model-framework to tflite or tensorflow. --model-version is optional and defaults to the current timestamp. If --model-name is already a model in your organization, the job adds a new version of it.

The SDKs can’t choose a framework: a job submitted from code always trains a TFLite model.

from viam.proto.app.mltraining import ModelType

# viam_client is a connected ViamClient
ml_training_client = viam_client.ml_training_client

job_id = await ml_training_client.submit_training_job(
    org_id="<org-id>",
    dataset_id="<dataset-id>",
    model_name="part-inspector",
    model_version="1",
    model_type=ModelType.MODEL_TYPE_SINGLE_LABEL_CLASSIFICATION,
    tags=["good-part", "defective-part"],
)
print(f"Training job: {job_id}")

The SDKs can’t choose a framework: a job submitted from code always trains a TFLite model.

// viamClient is a connected *app.ViamClient
mlTrainingClient := viamClient.MLTrainingClient()

jobID, err := mlTrainingClient.SubmitTrainingJob(
    ctx,
    app.SubmitTrainingJobArgs{
        DatasetID:      "<dataset-id>",
        OrganizationID: "<org-id>",
        ModelName:      "part-inspector",
        ModelVersion:   "1",
    },
    app.ModelTypeSingleLabelClassification,
    []string{"good-part", "defective-part"},
)
if err != nil {
    logger.Fatal(err)
}
fmt.Printf("Training job: %s\n", jobID)

The job starts, and you get its job ID. Training is billed by usage; see pricing.

2. Monitor the training job

In the Viam app, a job’s status is Preparing while it’s set up, Waiting for GPU while it’s queued for a cloud GPU, then Training, and it ends as Completed, Failed, or Canceled. The CLI and SDKs report Preparing as pending, and both Waiting for GPU and Training as in_progress. You get an email when it finishes.

  1. In the Viam app, go to DATA > MODELS and expand Active Training. It lists your organization’s training jobs with their status.
  2. Click a job ID to see its details and logs.

Check one job, and read its logs:

viam train get --job-id=<job-id>
viam train logs --job-id=<job-id>

List your organization’s training jobs:

viam train list --org-id=<org-id> --job-status=unspecified

job = await ml_training_client.get_training_job(id="<job-id>")
print(f"Status: {job.status}")

# List your organization's training jobs
for job in await ml_training_client.list_training_jobs(org_id="<org-id>"):
    print(f"{job.id}: {job.status}, model {job.model_name}")

job, err := mlTrainingClient.GetTrainingJob(ctx, "<job-id>")
if err != nil {
    logger.Fatal(err)
}
fmt.Printf("Status: %v\n", job.Status)

// List your organization's training jobs
jobs, err := mlTrainingClient.ListTrainingJobs(ctx, "<org-id>", app.TrainingStatusUnspecified)
if err != nil {
    logger.Fatal(err)
}
for _, j := range jobs {
    fmt.Printf("%s: %v, model %s\n", j.ID, j.Status, j.ModelName)
}

Training logs are kept for 7 days after the job finishes. Copy them before then if you need them longer.

3. Verify the trained model

When the job’s status is Completed, the model is in your organization’s registry. In the Viam app, go to DATA > MODELS and click the model’s name to open its details page, which lists its versions.

Built-in training holds back part of the dataset to validate the model, but it doesn’t report loss curves or test results. To see how the model does on images it hasn’t seen, evaluate it on a labeled test dataset before you deploy it.

Troubleshooting

The training job fails
  • Check the training logs. The error message usually says what went wrong.
  • The dataset is too small. Check it against the minimums in Prerequisites.
  • Fewer than two labels are selected for multi label classification. Multi label classification needs at least two labels.
  • Bounding boxes are out of range. For object detection, box coordinates must be between 0.0 and 1.0, with x_min less than x_max.
The model gives low confidence scores
  • Add more training images. Low confidence usually means the model hasn’t seen enough varied examples of each label.
  • Balance the labels. If one label dominates the dataset, the model may give low confidence to the others.
  • Check image quality. Blurry, dark, or low-resolution images make it harder for the model to learn.
  • Lower the confidence threshold. If the model is right but scores around 0.4 to 0.6, your threshold may be too high.
Training takes a long time
  • Large datasets take longer. A dataset with thousands of images can take an hour or more.
  • TensorFlow takes longer than TFLite. If training time matters, use TFLite.
  • The job is waiting for a GPU. A job stays Waiting for GPU until a cloud GPU is free.

What’s next