Skip to main content

Vision Tasks

Train models for image classification, object detection, and vision-language tasks.

Image Classification

Quick Start

aitraining image-classification \
  --model google/vit-base-patch16-224 \
  --data-path ./images/ \
  --project-name image-classifier

Parameters

ParameterDescriptionDefault
--modelBase modelgoogle/vit-base-patch16-224
--data-pathPath to image folderNone (required)
--project-nameOutput directoryproject-name
--image-columnImage column nameimage
--target-columnLabel column nametarget
--epochsTraining epochs3
--batch-sizeBatch size8
--lrLearning rate5e-5

Data Format

Folder structure:
images/
├── cat/
│   ├── cat1.jpg
│   └── cat2.jpg
├── dog/
│   ├── dog1.jpg
│   └── dog2.jpg
Or CSV with paths:
image,label
/path/to/cat1.jpg,cat
/path/to/dog1.jpg,dog

Example: Custom Classifier

aitraining image-classification \
  --model microsoft/resnet-50 \
  --data-path ./product_images/ \
  --project-name product-classifier \
  --epochs 10 \
  --batch-size 32 \
  --lr 5e-5

Image Regression

For predicting continuous values from images.

Quick Start

aitraining image-regression \
  --model google/vit-base-patch16-224 \
  --data-path ./images.csv \
  --image-column image \
  --target-column score \
  --project-name quality-scorer

Parameters

ParameterDescriptionDefault
--modelBase modelgoogle/vit-base-patch16-224
--data-pathPath to dataNone (required)
--project-nameOutput directoryproject-name
--image-columnImage column nameimage
--target-columnTarget value columntarget
--epochsTraining epochs3
--batch-sizeBatch size8
--lrLearning rate5e-5

Example: Age Prediction

aitraining image-regression \
  --model microsoft/resnet-50 \
  --data-path ./faces.csv \
  --image-column photo \
  --target-column age \
  --project-name age-predictor \
  --epochs 20

Object Detection

Quick Start

aitraining object-detection \
  --model facebook/detr-resnet-50 \
  --data-path ./coco_format/ \
  --project-name detector

Data Format

COCO-style format:
{
  "image": "image.jpg",
  "objects": {
    "bbox": [[x, y, width, height], ...],
    "categories": [0, 1, ...]
  }
}

Parameters

ParameterDescriptionDefault
--image-columnImage columnimage
--objects-columnObjects columnobjects
--image-square-sizeSquare size for images600

Example: Custom Detection

aitraining object-detection \
  --model facebook/detr-resnet-50 \
  --data-path ./annotations/ \
  --project-name custom-detector \
  --epochs 50 \
  --batch-size 8

Vision-Language Models (VLM)

Train models that understand both images and text.

Quick Start

aitraining vlm \
  --model google/paligemma-3b-pt-224 \
  --data-path ./image_captions.jsonl \
  --project-name vlm-model

Data Format

{
  "image": "image.jpg",
  "conversations": [
    {"role": "user", "content": "What's in this image?"},
    {"role": "assistant", "content": "A cat sitting on a couch."}
  ]
}

Parameters

ParameterDescriptionDefault
--modelBase modelgoogle/paligemma-3b-pt-224
--image-columnImage columnimage
--text-columnText/answer columntext
--prompt-text-columnPrompt/prefix columnprompt
--trainerTraining mode (vqa, captioning, segmentation, detection)vqa
--epochsTraining epochs3
--batch-sizeBatch size2
--lrLearning rate5e-5
--gradient-accumulationGradient accumulation steps4
--peftEnable LoRAFalse
--lora-rLoRA rank16
--lora-alphaLoRA alpha32

Example: Image Captioning

aitraining vlm \
  --model Qwen/Qwen2-VL-2B-Instruct \
  --data-path ./captions.jsonl \
  --project-name captioner \
  --epochs 3 \
  --batch-size 2 \
  --peft \
  --lora-r 16

Common Models

Image Classification

ModelParametersBest For
google/vit-base-patch16-22486MGeneral purpose
microsoft/resnet-5025MFast inference
facebook/convnext-base-22489MHigh accuracy

Object Detection

ModelParametersBest For
facebook/detr-resnet-5041MGeneral detection
facebook/detr-resnet-10160MHigher accuracy

Vision-Language

ModelParametersBest For
Qwen/Qwen2-VL-2B-Instruct2BBalanced
llava-hf/llava-1.5-7b-hf7BHigh quality

GPU Memory Tips

  • Use smaller batch sizes for large images
  • Enable gradient checkpointing for VLMs
  • Use LoRA for VLM training:
aitraining vlm \
  --model Qwen/Qwen2-VL-2B-Instruct \
  --data-path ./data.jsonl \
  --project-name vlm \
  --peft \
  --lora-r 16 \
  --batch-size 1 \
  --gradient-accumulation 8

Next Steps