Custom training scripts

Viam’s managed training handles TensorFlow and TFLite classification and detection out of the box. Custom training scripts are available for other use cases like a different framework, custom preprocessing, time-series data from a sequence dataset, or a training pipeline you want to share with your organization. Managed training doesn’t accept sequence datasets, so training on one always requires a custom training script.

Before writing your own, check the registry for existing training scripts and pre-trained models you can deploy directly. If a training script there fits your needs, skip ahead to Submit a training job.

How custom training scripts work

When you submit a custom training job, Viam:

1. Pulls your dataset and writes it into the container:
   • For a binary dataset, a JSONLines metadata file plus the image files.
   • For a sequence dataset, three Parquet files plus the image files.
2. Runs your script inside a Docker container with GPU access: either a Viam-managed container with a framework version you select, or a custom training container your organization registered.
3. Packages the artifacts your script writes to the output directory and publishes them as a new model version in your organization’s registry.

Because the script runs in a known container and receives standardized inputs, you can retrain on new datasets without changing the script or redeploying code.

Write a training script

A training script is a Python project that the Viam platform runs in the cloud. The platform provides your script with a dataset and an output directory; your script produces a trained model.

File structure

my-training/
├── model/
│   ├── training.py
│   └── __init__.py
└── setup.py

setup.py declares your dependencies:

from setuptools import find_packages, setup

setup(
    name="my-training",
    version="0.1",
    packages=find_packages(),
    include_package_data=True,
    install_requires=[
        # Add your dependencies here, for example:
        # "tensorflow>=2.11",
        # "numpy",
    ],
)

training.py

Your script always receives --model_output_directory. The arguments that point to the dataset depend on the dataset type. A job on a binary dataset passes these arguments:

A job on a sequence dataset passes three Parquet file paths instead of --dataset_file. See Sequence dataset inputs.

You can add custom arguments (like --num_epochs or --labels) and pass them when you submit the training job.

Here is the overall shape of a training script for a binary dataset:

import argparse
import json
import os


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset_file", dest="data_json",
                        type=str, required=True)
    parser.add_argument("--model_output_directory", dest="model_dir",
                        type=str, required=True)
    # Add custom arguments as needed:
    parser.add_argument("--num_epochs", dest="num_epochs",
                        type=int, default=200)
    return parser.parse_args()


def load_dataset(data_json):
    """Parse the JSONLines dataset file."""
    entries = []
    with open(data_json, "r") as f:
        for line in f:
            entries.append(json.loads(line))
    return entries


if __name__ == "__main__":
    args = parse_args()
    dataset = load_dataset(args.data_json)

    # --- Your training logic goes here ---
    # Use the dataset entries to train a model with
    # whatever framework fits your use case.
    model = ...
    labels = ...

    # Save model artifacts to the output directory.
    # The format must match what your ML model service expects.
    # For example, tflite_cpu expects a .tflite file and labels.txt.
    with open(os.path.join(args.model_dir, "model.tflite"), "wb") as f:
        f.write(model)
    with open(os.path.join(args.model_dir, "labels.txt"), "w") as f:
        f.write("\n".join(labels))

The critical parts:

• Parse arguments: Accept --model_output_directory plus the dataset arguments for your dataset type: --dataset_file for a binary dataset, or the three sequence file arguments for a sequence dataset.
• Read the dataset: For a binary dataset, each line in the JSONLines file is a JSON object with an image_path and either classification_annotations, bounding_box_annotations, or both. For a sequence dataset, read the three Parquet files instead.
• Save to the output directory: When the job completes, Viam packages everything in this directory and publishes it to the registry as a new model version. Files in a tmp/ subdirectory are excluded: use it for intermediate work.

If the script exits with a non-zero status or produces no files in the output directory, the training job is marked as failed.

Binary dataset file format

Each line of the --dataset_file file is a JSON object like this:

{
  "image_path": "/path/to/data/img1.jpeg",
  "classification_annotations": [{ "annotation_label": "blue_star" }],
  "bounding_box_annotations": [
    {
      "annotation_label": "blue_star",
      "x_min_normalized": 0.382,
      "x_max_normalized": 0.51,
      "y_min_normalized": 0.356,
      "y_max_normalized": 0.527
    }
  ]
}

Bounding box coordinates are normalized to the range 0.0-1.0 relative to image dimensions. For classification, read classification_annotations. For object detection, read bounding_box_annotations. See the example training script for complete parsing functions that handle both annotation types.

Sequence dataset inputs

A job on a sequence dataset doesn’t get a --dataset_file. Instead, your script receives three Parquet files: --binary_data_file, --tabular_data_file, and --sequences_file. It still receives --model_output_directory. The images are files on disk, and the Parquet files hold their metadata and your readings. Join the three files on sequence_id. See Sequence dataset format for the columns. To record and export a sequence dataset to test your script on, follow the sequences tutorial.

Accessing Viam APIs

The platform provides API_KEY and API_KEY_ID environment variables if your script needs to call Viam APIs during training, for example, to query additional data:

import os
from viam.rpc.dial import DialOptions
from viam.app.viam_client import ViamClient


async def connect() -> ViamClient:
    dial_options = DialOptions.with_api_key(
        os.environ.get("API_KEY"), os.environ.get("API_KEY_ID")
    )
    return await ViamClient.create_from_dial_options(dial_options)

Match your output to an ML model service

After training, Viam packages your output directory and publishes it as a model version in your organization’s registry. To run the model on a machine, the ML model service you deploy with it must be able to load the files your script wrote.

For TFLite models loaded with tflite_cpu, save a model.tflite file and a labels.txt listing each label on its own line. This is what the classification-tflite example produces.

For other frameworks, consult the module’s README on the registry – for example, tensorflow-cpu, onnx-cpu, or torch-cpu – for the file format and tensor shape expectations. Write your output in that format so the deployed model loads correctly.

After uploading, follow Deploy a model to a machine to add the matching ML model service and vision service to your machine.

Complete example

For a full working training script, see the classification-tflite example on GitHub. It trains a TFLite single-label classification model using TensorFlow and Keras.

Test and upload

Before submitting a cloud training job, test your script locally against an exported dataset.

Export a dataset

viam dataset export --destination=<destination> --dataset-id=<dataset-id>

For a binary dataset, this downloads the image files and a dataset.jsonl metadata file. To download only the JSONL file without the images, add --only-jsonl.

For a sequence dataset, this downloads a zip of three Parquet files plus the images. To download only the zip without the images, add --only-parquet. See Export a sequence dataset.

You can get the dataset ID from the DATASETS tab or by running viam dataset list.

Test locally with Docker

The test-local command runs your training script inside the same Docker container that cloud training uses. This catches problems that plain Python testing misses: missing system dependencies, Python version differences, and package conflicts. Testing locally is free, while a cloud training job bills GPU time from the moment it starts, so it’s a cheap way to find a broken script.

test-local runs the container without a GPU, so your script runs on the CPU. Use a small dataset or a low epoch count, and make sure your script can fall back to the CPU. It’s a quick check that the script runs, not a substitute for the cloud training job.

test-local passes your script only --dataset_file, so it works only with scripts for binary datasets. To test a script for a sequence dataset, run it directly.

viam training-script test-local \
  --training-script-directory=my-training/ \
  --dataset-file=dataset.jsonl \
  --dataset-root=<destination> \
  --container-id=<container-id> \
  --model-output-directory=<output-dir>

The --dataset-file path is relative to --dataset-root. The command mounts your script, dataset, and output directories into the container.

The --container-id flag selects the container to run in, either a Viam-managed container or a custom training container, so you test in the same environment the cloud training job uses. To find container IDs, see Submit a training job.

You can pass custom arguments with --custom-args:

viam training-script test-local \
  --training-script-directory=my-training/ \
  --dataset-file=dataset.jsonl \
  --dataset-root=<destination> \
  --container-id=<container-id> \
  --model-output-directory=<output-dir> \
  --custom-args=num_epochs=5,labels="label1 label2"

Run without Docker

For a quick check without Docker, run your script directly. For a binary dataset:

python3 -m model.training --dataset_file=<path/to/dataset.jsonl> \
    --model_output_directory=<output-dir>

For a sequence dataset, unzip the export first, then pass the three Parquet files:

python3 -m model.training \
    --binary_data_file=<unzipped-dir>/binary_data.parquet \
    --tabular_data_file=<unzipped-dir>/tabular_data.parquet \
    --sequences_file=<unzipped-dir>/sequences.parquet \
    --model_output_directory=<output-dir>

In an export, the path column in binary_data.parquet is relative to the export directory. In a cloud training job, path is an absolute path. Resolve relative paths against the export directory so the same script works in both places.

Package and upload

tar -czvf my-training.tar.gz my-training/
viam training-script upload --path=my-training.tar.gz \
  --org-id=<org-id> --script-name=my-training-script

You can also specify --framework, --type, --visibility, and --description when uploading. Scripts default to private (visible only within your organization). Set --visibility=public to share with other organizations through the registry. See the CLI reference for the full list of flags.

To find your organization ID, run viam organization list.

After uploading, your script appears in the registry.

Submit a training job

Once a training script is in the registry, whether you uploaded it or are using someone else’s, submit a training job to run it against a dataset.

Every custom training job runs inside a container on GPU-backed cloud infrastructure. Viam-managed containers are based on Python 3.10 with TensorFlow pre-installed. The currently supported versions are tf:2.16, tf:2.17, and tf:2.21. To use a different framework (PyTorch, scikit-learn, or anything installable with pip), add it to the install_requires list in setup.py; Viam installs the packages in the container before running your script.

If your dependencies take a long time to install, build them into a custom training container instead so each job starts training right away.

Custom training jobs run on premium GPUs and are billed at a higher rate than managed training jobs; see pricing for plan details.

To list available containers with their IDs, framework versions, and end-of-life dates, including your organization’s custom containers, run:

viam train containers list --org-id=<org-id>

viam train submit custom from-registry --dataset-id=<dataset-id> \
  --org-id=<org-id> --model-name=my-model \
  --model-version=1 --version=1 \
  --script-name=<namespace>:<script-name> \
  --container-id=<container-id> \
  --args=num_epochs=100,labels="'label1 label2'"

Monitor and debug

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

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

Training logs expire after 7 days. You will receive an email when your training job completes.

If a job fails, check the logs first: the error message usually indicates the problem. Note that training scripts may emit log lines at the error level and still succeed; check the final job status rather than individual log lines.