cog/keypoint — Keypoint training, preview, export

Train a generic keypoint detector, preview augmented data, run a dry-run, or export the latest checkpoint to ONNX.

Usage

# Preview a batch of augmented patches
s6 cog keypoint --config ./cfg/ds.json --preview-data

# Dry‑run (log graph and run one training step)
s6 cog keypoint --config ./cfg/ds.json --dry-run

# Train for 50 epochs and export an ONNX model
s6 cog keypoint --config ./cfg/ds.json --train -e 50 -b 16 -lr 1e-3 \
  --deploy model.onnx

# Restore latest checkpoint before continuing training
s6 cog keypoint --config ./cfg/ds.json --train --restore latest

How it works

  • Loads dataset/augmentation from a JSON config via s6.app.cog.augmented_dataset.AugmentedKeypointDataset.from_json.

  • Uses s6.nn.backbones.keypoints.GenericKeypointModel(in_channels=1) with Smooth L1 loss on normalized keypoints in [0, 1).

  • Logs to TensorBoard under logs/cog unless --no-tb is set.

  • Saves checkpoints under checkpoints/ and optionally exports ONNX with outputs named keypoints and heatmaps.

Key flags

  • --config — path to dataset config JSON (required)

  • -t/--train — run training; --dry-run for one step + graph

  • -e/--epochs, -b/--batch_size, -lr/--learning_rate

  • --deploy — ONNX filename to export after training

  • --restore — checkpoint path or latest to resume

  • --no-tb — disable TensorBoard logging; --log-dir to change root