Training dataset formats

A custom training script receives its dataset as files on disk. What those files are depends on the dataset type: a binary dataset arrives as one JSONLines file plus the files it lists, and a sequence dataset arrives as three Parquet files plus the images. viam dataset export writes the same formats, so you can test a script on an export before you run it in the cloud.

Training script arguments

A custom training job runs your script as python3 -m model.training with these arguments:

A job on a binary dataset gets --dataset_file and none of the three Parquet file arguments. A job on a sequence dataset gets the three Parquet file arguments and no --dataset_file.

You can add your own arguments, such as --num_epochs, and pass values for them when you submit the job. Viam passes each key and value as --<key>=<value>. Your own arguments can’t reuse the names in the table above.

Binary datasets

Export layout

Running viam dataset export on a binary dataset writes:

<destination>/
  dataset.jsonl
  data/
    <file>

To write only dataset.jsonl, without the files, add --only-jsonl.

JSONLines file

Each line of the file is a JSON object describing one file in the dataset:

{
  "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
    }
  ]
}

An exported file’s lines also carry metadata about each file, such as timestamp, binary_data_id, part_id, and component_name.

In an export, image_path starts with the --destination you passed to viam dataset export. If that was a relative path, it’s relative to the directory you ran the command from.

For parsing functions that handle both annotation types, see the classification-tflite example script.

Sequence datasets

Export layout

Running viam dataset export on a sequence dataset writes:

<destination>/
  <dataset-id>.zip
    binary_data.parquet
    tabular_data.parquet
    sequences.parquet
  binary_data/
    <binary-data-id><extension>

Binary data IDs contain slashes, so the image files sit in nested folders under binary_data/. To download only the zip, without the images, add --only-parquet.

In an exported binary_data.parquet, the path column is relative to the export directory. In a training job, path is an absolute path, starting with /gcs/, that your script can open directly. To produce an export you can inspect, follow the sequences tutorial.

binary_data.parquet

One row per image in each sequence.

tabular_data.parquet

One row per reading in each sequence.

sequences.parquet

One row per sequence.

The API and SDKs call these fields start_time and end_time.

Join the files

Join binary_data.parquet and tabular_data.parquet to sequences.parquet on sequence_id.

The export doesn’t resample or align the data. Each resource keeps its own sample rate, so a script that needs images and readings at the same moments has to match them itself, for example by nearest timestamp. All timestamps are in UTC, and libraries such as pandas load them as timezone-aware UTC values.

Training job environment

Output directory

When the job finishes, Viam packages every file in --model_output_directory and publishes it to your organization’s registry as a new model version. Files in a tmp/ subdirectory are left out, so use it for checkpoints and other intermediate files.

The job fails if your script exits with a non-zero status or writes no files to the output directory.

Viam API credentials

The job sets the API_KEY and API_KEY_ID environment variables, so your script can call Viam APIs during training, for example to query more 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)

See Train a model with a custom training script.