# SkillPatch skill: huggingface-vision-trainer A comprehensive skill for training and fine-tuning vision models on Hugging Face Jobs cloud GPUs. Covers object detection (D-FINE, RT-DETR v2, DETR, YOLOS), image classification (timm models like MobileNetV3, MobileViT, ResNet, ViT/DINOv3), and SAM/SAM2 segmentation. Includes dataset preparation, augmentation, evaluation metrics, hardware selection, cost estimation, and Hub persistence. You (the agent) were given this URL and asked to use this skill. This is a **public** skill — no SkillPatch account, API key, or prior setup is required. Two ways to use it: **1. Use it right now** — the skill's complete file tree (SKILL.md + all reference files) is inlined below; read `SKILL.md` first, then follow it, consulting the other files as it directs. **2. Install the exact package onto disk** (recommended if you can run a shell — this reproduces the skill byte-for-byte, including any binary assets that can't be inlined): ```bash mkdir -p .claude/skills/huggingface-vision-trainer curl -sSL https://skillpatch.dev/install_skill/huggingface-vision-trainer | tar -xz -C .claude/skills/ ``` (`.claude/skills/` is Claude Code's convention; use whatever directory your agent loads skills from.) --- ## Skill files (12) - `SKILL.md` - `references/finetune_sam2_trainer.md` - `references/hub_saving.md` - `references/image_classification_training_notebook.md` - `references/object_detection_training_notebook.md` - `references/reliability_principles.md` - `references/timm_trainer.md` - `scripts/dataset_inspector.py` - `scripts/estimate_cost.py` - `scripts/image_classification_training.py` - `scripts/object_detection_training.py` - `scripts/sam_segmentation_training.py` ### `SKILL.md` ````markdown --- name: huggingface-vision-trainer description: Trains and fine-tunes vision models for object detection (D-FINE, RT-DETR v2, DETR, YOLOS), image classification (timm models — MobileNetV3, MobileViT, ResNet, ViT/DINOv3 — plus any Transformers classifier), and SAM/SAM2 segmentation using Hugging Face Transformers on Hugging Face Jobs cloud GPUs. Covers COCO-format dataset preparation, Albumentations augmentation, mAP/mAR evaluation, accuracy metrics, SAM segmentation with bbox/point prompts, DiceCE loss, hardware selection, cost estimation, Trackio monitoring, and Hub persistence. Use when users mention training object detection, image classification, SAM, SAM2, segmentation, image matting, DETR, D-FINE, RT-DETR, ViT, timm, MobileNet, ResNet, bounding box models, or fine-tuning vision models on Hugging Face Jobs. --- # Vision Model Training on Hugging Face Jobs Train object detection, image classification, and SAM/SAM2 segmentation models on managed cloud GPUs. No local GPU setup required—results are automatically saved to the Hugging Face Hub. ## When to Use This Skill Use this skill when users want to: - Fine-tune object detection models (D-FINE, RT-DETR v2, DETR, YOLOS) on cloud GPUs or local - Fine-tune image classification models (timm: MobileNetV3, MobileViT, ResNet, ViT/DINOv3, or any Transformers classifier) on cloud GPUs or local - Fine-tune SAM or SAM2 models for segmentation / image matting using bbox or point prompts - Train bounding-box detectors on custom datasets - Train image classifiers on custom datasets - Train segmentation models on custom mask datasets with prompts - Run vision training jobs on Hugging Face Jobs infrastructure - Ensure trained vision models are permanently saved to the Hub ## Related Skills - **`hugging-face-jobs`** — General HF Jobs infrastructure: token authentication, hardware flavors, timeout management, cost estimation, secrets, environment variables, scheduled jobs, and result persistence. **Refer to the Jobs skill for any non-training-specific Jobs questions** (e.g., "how do secrets work?", "what hardware is available?", "how do I pass tokens?"). - **`hugging-face-model-trainer`** — TRL-based language model training (SFT, DPO, GRPO). Use that skill for text/language model fine-tuning. ## Local Script Execution Helper scripts use PEP 723 inline dependencies. Run them with `uv run`: ```bash uv run scripts/dataset_inspector.py --dataset username/dataset-name --split train uv run scripts/estimate_cost.py --help ``` ## Prerequisites Checklist Before starting any training job, verify: ### Account & Authentication - Hugging Face Account with [Pro](https://hf.co/pro), [Team](https://hf.co/enterprise), or [Enterprise](https://hf.co/enterprise) plan (Jobs require paid plan) - Authenticated login: Check with `hf_whoami()` (tool) or `hf auth whoami` (terminal) - Token has **write** permissions - **MUST pass token in job secrets** — see directive #3 below for syntax (MCP tool vs Python API) ### Dataset Requirements — Object Detection - Dataset must exist on Hub - Annotations must use the `objects` column with `bbox`, `category` (and optionally `area`) sub-fields - Bboxes can be in **xywh (COCO)** or **xyxy (Pascal VOC)** format — auto-detected and converted - Categories can be **integers or strings** — strings are auto-remapped to integer IDs - `image_id` column is **optional** — generated automatically if missing - **ALWAYS validate unknown datasets** before GPU training (see Dataset Validation section) ### Dataset Requirements — Image Classification - Dataset must exist on Hub - Must have an **`image` column** (PIL images) and a **`label` column** (integer class IDs or strings) - The label column can be `ClassLabel` type (with names) or plain integers/strings — strings are auto-remapped - Common column names auto-detected: `label`, `labels`, `class`, `fine_label` - **ALWAYS validate unknown datasets** before GPU training (see Dataset Validation section) ### Dataset Requirements — SAM/SAM2 Segmentation - Dataset must exist on Hub - Must have an **`image` column** (PIL images) and a **`mask` column** (binary ground-truth segmentation mask) - Must have a **prompt** — either: - A **`prompt` column** with JSON containing `{"bbox": [x0,y0,x1,y1]}` or `{"point": [x,y]}` - OR a dedicated **`bbox`** column with `[x0,y0,x1,y1]` values - OR a dedicated **`point`** column with `[x,y]` or `[[x,y],...]` values - Bboxes should be in **xyxy** format (absolute pixel coordinates) - Example dataset: `merve/MicroMat-mini` (image matting with bbox prompts) - **ALWAYS validate unknown datasets** before GPU training (see Dataset Validation section) ### Critical Settings - **Timeout must exceed expected training time** — Default 30min is TOO SHORT. See directive #6 for recommended values. - **Hub push must be enabled** — `push_to_hub=True`, `hub_model_id="username/model-name"`, token in `secrets` ## Dataset Validation **Validate dataset format BEFORE launching GPU training to prevent the #1 cause of training failures: format mismatches.** **ALWAYS validate for** unknown/custom datasets or any dataset you haven't trained with before. **Skip for** `cppe-5` (the default in the training script). ### Running the Inspector **Option 1: Via HF Jobs (recommended — avoids local SSL/dependency issues):** ```python hf_jobs("uv", { "script": "path/to/dataset_inspector.py", "script_args": ["--dataset", "username/dataset-name", "--split", "train"] }) ``` **Option 2: Locally:** ```bash uv run scripts/dataset_inspector.py --dataset username/dataset-name --split train ``` **Option 3: Via `HfApi().run_uv_job()` (if hf_jobs MCP unavailable):** ```python from huggingface_hub import HfApi api = HfApi() api.run_uv_job( script="scripts/dataset_inspector.py", script_args=["--dataset", "username/dataset-name", "--split", "train"], flavor="cpu-basic", timeout=300, ) ``` ### Reading Results - **`✓ READY`** — Dataset is compatible, use directly - **`✗ NEEDS FORMATTING`** — Needs preprocessing (mapping code provided in output) ## Automatic Bbox Preprocessing The object detection training script (`scripts/object_detection_training.py`) automatically handles bbox format detection (xyxy→xywh conversion), bbox sanitization, `image_id` generation, string category→integer remapping, and dataset truncation. **No manual preprocessing needed** — just ensure the dataset has `objects.bbox` and `objects.category` columns. ## Training workflow Copy this checklist and track progress: ``` Training Progress: - [ ] Step 1: Verify prerequisites (account, token, dataset) - [ ] Step 2: Validate dataset format (run dataset_inspector.py) - [ ] Step 3: Ask user about dataset size and validation split - [ ] Step 4: Prepare training script (OD: scripts/object_detection_training.py, IC: scripts/image_classification_training.py, SAM: scripts/sam_segmentation_training.py) - [ ] Step 5: Save script locally, submit job, and report details ``` **Step 1: Verify prerequisites** Follow the Prerequisites Checklist above. **Step 2: Validate dataset** Run the dataset inspector BEFORE spending GPU time. See "Dataset Validation" section above. **Step 3: Ask user preferences** ALWAYS use the AskUserQuestion tool with option-style format: ```python AskUserQuestion({ "questions": [ { "question": "Do you want to run a quick test with a subset of the data first?", "header": "Dataset Size", "options": [ {"label": "Quick test run (10% of data)", "description": "Faster, cheaper (~30-60 min, ~$2-5) to validate setup"}, {"label": "Full dataset (Recommended)", "description": "Complete training for best model quality"} ], "multiSelect": false }, { "question": "Do you want to create a validation split from the training data?", "header": "Split data", "options": [ {"label": "Yes (Recommended)", "description": "Automatically split 15% of training data for validation"}, {"label": "No", "description": "Use existing validation split from dataset"} ], "multiSelect": false }, { "question": "Which GPU hardware do you want to use?", "header": "Hardware Flavor", "options": [ {"label": "t4-small ($0.40/hr)", "description": "1x T4, 16 GB VRAM — sufficient for all OD models under 100M params"}, {"label": "l4x1 ($0.80/hr)", "description": "1x L4, 24 GB VRAM — more headroom for large images or batch sizes"}, {"label": "a10g-large ($1.50/hr)", "description": "1x A10G, 24 GB VRAM — faster training, more CPU/RAM"}, {"label": "a100-large ($2.50/hr)", "description": "1x A100, 80 GB VRAM — fastest, for very large datasets or image sizes"} ], "multiSelect": false } ] }) ``` **Step 4: Prepare training script** For object detection, use [scripts/object_detection_training.py](scripts/object_detection_training.py) as the production-ready template. For image classification, use [scripts/image_classification_training.py](scripts/image_classification_training.py). For SAM/SAM2 segmentation, use [scripts/sam_segmentation_training.py](scripts/sam_segmentation_training.py). All scripts use `HfArgumentParser` — all configuration is passed via CLI arguments in `script_args`, NOT by editing Python variables. For timm model details, see [references/timm_trainer.md](references/timm_trainer.md). For SAM2 training details, see [references/finetune_sam2_trainer.md](references/finetune_sam2_trainer.md). **Step 5: Save script, submit job, and report** 1. **Save the script locally** to `submitted_jobs/` in the workspace root (create if needed) with a descriptive name like `training__.py`. Tell the user the path. 2. **Submit** using `hf_jobs` MCP tool (preferred) or `HfApi().run_uv_job()` — see directive #1 for both methods. Pass all config via `script_args`. 3. **Report** the job ID (from `.id` attribute), monitoring URL, Trackio dashboard (`https://huggingface.co/spaces/{username}/trackio`), expected time, and estimated cost. 4. **Wait for user** to request status checks — don't poll automatically. Training jobs run asynchronously and can take hours. ## Critical directives These rules prevent common failures. Follow them exactly. ### 1. Job submission: `hf_jobs` MCP tool vs Python API **`hf_jobs()` is an MCP tool, NOT a Python function.** Do NOT try to import it from `huggingface_hub`. Call it as a tool: ``` hf_jobs("uv", {"script": training_script_content, "flavor": "a10g-large", "timeout": "4h", "secrets": {"HF_TOKEN": "$HF_TOKEN"}}) ``` **If `hf_jobs` MCP tool is unavailable**, use the Python API directly: ```python from huggingface_hub import HfApi, get_token api = HfApi() job_info = api.run_uv_job( script="path/to/training_script.py", # file PATH, NOT content script_args=["--dataset_name", "cppe-5", ...], flavor="a10g-large", timeout=14400, # seconds (4 hours) env={"PYTHONUNBUFFERED": "1"}, secrets={"HF_TOKEN": get_token()}, # MUST use get_token(), NOT "$HF_TOKEN" ) print(f"Job ID: {job_info.id}") ``` **Critical differences between the two methods:** | | `hf_jobs` MCP tool | `HfApi().run_uv_job()` | |---|---|---| | `script` param | Python code string or URL (NOT local paths) | File path to `.py` file (NOT content) | | Token in secrets | `"$HF_TOKEN"` (auto-replaced) | `get_token()` (actual token value) | | Timeout format | String (`"4h"`) | Seconds (`14400`) | **Rules for both methods:** - The training script MUST include PEP 723 inline metadata with dependencies - Do NOT use `image` or `command` parameters (those belong to `run_job()`, not `run_uv_job()`) ### 2. Authentication via job secrets + explicit hub_token injection **Job config** MUST include the token in secrets — syntax depends on submission method (see table above). **Training script requirement:** The Transformers `Trainer` calls `create_repo(token=self.args.hub_token)` during `__init__()` when `push_to_hub=True`. The training script MUST inject `HF_TOKEN` into `training_args.hub_token` AFTER parsing args but BEFORE creating the `Trainer`. The template `scripts/object_detection_training.py` already includes this: ```python hf_token = os.environ.get("HF_TOKEN") if training_args.push_to_hub and not training_args.hub_token: if hf_token: training_args.hub_token = hf_token ``` If you write a custom script, you MUST include this token injection before the `Trainer(...)` call. - Do NOT call `login()` in custom scripts unless replicating the full pattern from `scripts/object_detection_training.py` - Do NOT rely on implicit token resolution (`hub_token=None`) — unreliable in Jobs - See the `hugging-face-jobs` skill → *Token Usage Guide* for full details ### 3. JobInfo attribute Access the job identifier using `.id` (NOT `.job_id` or `.name` — these don't exist): ```python job_info = api.run_uv_job(...) # or hf_jobs("uv", {...}) job_id = job_info.id # Correct -- returns string like "687fb701029421ae5549d998" ``` ### 4. Required training flags and HfArgumentParser boolean syntax `scripts/object_detection_training.py` uses `HfArgumentParser` — all config is passed via `script_args`. Boolean arguments have two syntaxes: - **`bool` fields** (e.g., `push_to_hub`, `do_train`): Use as bare flags (`--push_to_hub`) or negate with `--no_` prefix (`--no_remove_unused_columns`) - **`Optional[bool]` fields** (e.g., `greater_is_better`): MUST pass explicit value (`--greater_is_better True`). Bare `--greater_is_better` causes `error: expected one argument` Required flags for object detection: ``` --no_remove_unused_columns # MUST: preserves image column for pixel_values --no_eval_do_concat_batches # MUST: images have different numbers of target boxes --push_to_hub # MUST: environment is ephemeral --hub_model_id username/model-name --metric_for_best_model eval_map --greater_is_better True # MUST pass "True" explicitly (Optional[bool]) --do_train --do_eval ``` Required flags for image classification: ``` --no_remove_unused_columns # MUST: preserves image column for pixel_values --push_to_hub # MUST: environment is ephemeral --hub_model_id username/model-name --metric_for_best_model eval_accuracy --greater_is_better True # MUST pass "True" explicitly (Optional[bool]) --do_train --do_eval ``` Required flags for SAM/SAM2 segmentation: ``` --remove_unused_columns False # MUST: preserves input_boxes/input_points --push_to_hub # MUST: environment is ephemeral --hub_model_id username/model-name --do_train --prompt_type bbox # or "point" --dataloader_pin_memory False # MUST: avoids pin_memory issues with custom collator ``` ### 5. Timeout management Default 30 min is TOO SHORT for object detection. Set minimum 2-4 hours. Add 30% buffer for model loading, preprocessing, and Hub push. | Scenario | Timeout | |----------|---------| | Quick test (100-200 images, 5-10 epochs) | 1h | | Development (500-1K images, 15-20 epochs) | 2-3h | | Production (1K-5K images, 30 epochs) | 4-6h | | Large dataset (5K+ images) | 6-12h | ### 6. Trackio monitoring Trackio is **always enabled** in the object detection training script — it calls `trackio.init()` and `trackio.finish()` automatically. No need to pass `--report_to trackio`. The project name is taken from `--output_dir` and the run name from `--run_name`. For image classification, pass `--report_to trackio` in `TrainingArguments`. Dashboard at: `https://huggingface.co/spaces/{username}/trackio` ## Model & hardware selection ### Recommended object detection models | Model | Params | Use case | |-------|--------|----------| | `ustc-community/dfine-small-coco` | 10.4M | Best starting point — fast, cheap, SOTA quality | | `PekingU/rtdetr_v2_r18vd` | 20.2M | Lightweight real-time detector | | `ustc-community/dfine-large-coco` | 31.4M | Higher accuracy, still efficient | | `PekingU/rtdetr_v2_r50vd` | 43M | Strong real-time baseline | | `ustc-community/dfine-xlarge-obj365` | 63.5M | Best accuracy (pretrained on Objects365) | | `PekingU/rtdetr_v2_r101vd` | 76M | Largest RT-DETR v2 variant | Start with `ustc-community/dfine-small-coco` for fast iteration. Move to D-FINE Large or RT-DETR v2 R50 for better accuracy. ### Recommended image classification models All `timm/` models work out of the box via `AutoModelForImageClassification` (loaded as `TimmWrapperForImageClassification`). See [references/timm_trainer.md](references/timm_trainer.md) for details. | Model | Params | Use case | |-------|--------|----------| | `timm/mobilenetv3_small_100.lamb_in1k` | 2.5M | Ultra-lightweight — mobile/edge, fastest training | | `timm/mobilevit_s.cvnets_in1k` | 5.6M | Mobile transformer — good accuracy/speed trade-off | | `timm/resnet50.a1_in1k` | 25.6M | Strong CNN baseline — reliable, well-studied | | `timm/vit_base_patch16_dinov3.lvd1689m` | 86.6M | Best accuracy — DINOv3 self-supervised ViT | Start with `timm/mobilenetv3_small_100.lamb_in1k` for fast iteration. Move to `timm/resnet50.a1_in1k` or `timm/vit_base_patch16_dinov3.lvd1689m` for better accuracy. ### Recommended SAM/SAM2 segmentation models | Model | Params | Use case | |-------|--------|----------| | `facebook/sam2.1-hiera-tiny` | 38.9M | Fastest SAM2 — good for quick experiments | | `facebook/sam2.1-hiera-small` | 46.0M | Best starting point — good quality/speed balance | | `facebook/sam2.1-hiera-base-plus` | 80.8M | Higher capacity for complex segmentation | | `facebook/sam2.1-hiera-large` | 224.4M | Best SAM2 accuracy — requires more VRAM | | `facebook/sam-vit-base` | 93.7M | Original SAM — ViT-B backbone | | `facebook/sam-vit-large` | 312.3M | Original SAM — ViT-L backbone | | `facebook/sam-vit-huge` | 641.1M | Original SAM — ViT-H, best SAM v1 accuracy | Start with `facebook/sam2.1-hiera-small` for fast iteration. SAM2 models are generally more efficient than SAM v1 at similar quality. Only the mask decoder is trained by default (vision and prompt encoders are frozen). ### Hardware recommendation All recommended OD and IC models are under 100M params — **`t4-small` (16 GB VRAM, $0.40/hr) is sufficient for all of them.** Image classification models are generally smaller and faster than object detection models — `t4-small` handles even ViT-Base comfortably. For SAM2 models up to `hiera-base-plus`, `t4-small` is sufficient since only the mask decoder is trained. For `sam2.1-hiera-large` or SAM v1 models, use `l4x1` or `a10g-large`. Only upgrade if you hit OOM from large batch sizes — reduce batch size first before switching hardware. Common upgrade path: `t4-small` → `l4x1` ($0.80/hr, 24 GB) → `a10g-large` ($1.50/hr, 24 GB). For full hardware flavor list: refer to the `hugging-face-jobs` skill. For cost estimation: run `scripts/estimate_cost.py`. ## Quick start — Object Detection The `script_args` below are the same for both submission methods. See directive #1 for the critical differences between them. ```python OD_SCRIPT_ARGS = [ "--model_name_or_path", "ustc-community/dfine-small-coco", "--dataset_name", "cppe-5", "--image_square_size", "640", "--output_dir", "dfine_finetuned", "--num_train_epochs", "30", "--per_device_train_batch_size", "8", "--learning_rate", "5e-5", "--eval_strategy", "epoch", "--save_strategy", "epoch", "--save_total_limit", "2", "--load_best_model_at_end", "--metric_for_best_model", "eval_map", "--greater_is_better", "True", "--no_remove_unused_columns", "--no_eval_do_concat_batches", "--push_to_hub", "--hub_model_id", "username/mode ... ```` ### `references/finetune_sam2_trainer.md` ````markdown # Fine-tuning SAM2 with HF Trainer Fine-tune SAM2.1 on a small part of the MicroMat dataset for image matting, using the Hugging Face Trainer with a custom loss function. ```python !pip install -q transformers datasets monai trackio ``` ## Load and explore the dataset ```python from datasets import load_dataset dataset = load_dataset("merve/MicroMat-mini", split="train") dataset ``` ```python dataset = dataset.train_test_split(test_size=0.1) train_ds = dataset["train"] val_ds = dataset["test"] ``` ```python import json train_ds[0] ``` ```python json.loads(train_ds["prompt"][0])["bbox"] ``` ## Visualize a sample ```python import matplotlib.pyplot as plt import numpy as np def show_mask(mask, ax, bbox): color = np.array([0.12, 0.56, 1.0, 0.6]) mask = np.array(mask) h, w = mask.shape mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, 4) ax.imshow(mask_image) x0, y0, x1, y1 = bbox ax.add_patch( plt.Rectangle( (x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor="lime", linewidth=2 ) ) example = train_ds[0] image = np.array(example["image"]) ground_truth_mask = np.array(example["mask"]) fig, ax = plt.subplots() ax.imshow(image) show_mask(ground_truth_mask, ax, json.loads(example["prompt"])["bbox"]) ax.set_title("Ground truth mask") ax.set_axis_off() plt.show() ``` ## Build the dataset and collator `SAMDataset` wraps each sample into the format expected by the SAM2 processor. Ground-truth masks are stored under the key `"labels"` so the Trainer automatically pops them before calling `model.forward()`. ```python from torch.utils.data import Dataset import torch import torch.nn.functional as F class SAMDataset(Dataset): def __init__(self, dataset, processor): self.dataset = dataset self.processor = processor def __len__(self): return len(self.dataset) def __getitem__(self, idx): item = self.dataset[idx] image = item["image"] prompt = json.loads(item["prompt"])["bbox"] inputs = self.processor(image, input_boxes=[[prompt]], return_tensors="pt") inputs["labels"] = (np.array(item["mask"]) > 0).astype(np.float32) inputs["original_image_size"] = torch.tensor(image.size[::-1]) return inputs def collate_fn(batch): pixel_values = torch.cat([item["pixel_values"] for item in batch], dim=0) original_sizes = torch.stack([item["original_sizes"] for item in batch]) input_boxes = torch.cat([item["input_boxes"] for item in batch], dim=0) labels = torch.cat( [ F.interpolate( torch.as_tensor(x["labels"]).unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode="nearest", ) for x in batch ], dim=0, ).long() return { "pixel_values": pixel_values, "original_sizes": original_sizes, "input_boxes": input_boxes, "labels": labels, "original_image_size": torch.stack( [item["original_image_size"] for item in batch] ), "multimask_output": False, } ``` ```python from transformers import Sam2Processor processor = Sam2Processor.from_pretrained("facebook/sam2.1-hiera-small") train_dataset = SAMDataset(dataset=train_ds, processor=processor) val_dataset = SAMDataset(dataset=val_ds, processor=processor) ``` ## Load model and freeze encoder layers ```python from transformers import Sam2Model model = Sam2Model.from_pretrained("facebook/sam2.1-hiera-small") for name, param in model.named_parameters(): if name.startswith("vision_encoder") or name.startswith("prompt_encoder"): param.requires_grad_(False) ``` ## Inference before training ```python item = val_ds[1] img = item["image"] bbox = json.loads(item["prompt"])["bbox"] inputs = processor(images=img, input_boxes=[[bbox]], return_tensors="pt").to( model.device ) with torch.no_grad(): outputs = model(**inputs) masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0] preds = masks.squeeze(0) mask = (preds[0] > 0).cpu().numpy() overlay = np.asarray(img, dtype=np.uint8).copy() overlay[mask] = 0.55 * overlay[mask] + 0.45 * np.array([0, 255, 0], dtype=np.float32) plt.imshow(overlay) plt.title("Before training") plt.axis("off") plt.show() ``` ## Define custom loss SAM2 does not compute loss in its `forward()`, so we provide a `compute_loss_func` to the Trainer. The Trainer pops `"labels"` from the batch before calling `model(**inputs)`, then passes `(outputs, labels)` to this function. ```python import monai from transformers import Trainer, TrainingArguments import trackio seg_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction="mean") def compute_loss(outputs, labels, num_items_in_batch=None): predicted_masks = outputs.pred_masks.squeeze(1) return seg_loss(predicted_masks, labels.float()) ``` ## Train with Trainer Key settings: - `remove_unused_columns=False`: the Trainer must keep `input_boxes`, `original_sizes`, etc. that are not in the model's `forward()` signature. - `compute_loss_func`: our custom DiceCE loss. - `report_to="trackio"`: logs the training loss to trackio. ```python training_args = TrainingArguments( output_dir="sam2-finetuned", num_train_epochs=30, per_device_train_batch_size=4, learning_rate=1e-5, weight_decay=0, logging_steps=1, save_strategy="epoch", save_total_limit=2, remove_unused_columns=False, dataloader_pin_memory=False, report_to="trackio", ) trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, data_collator=collate_fn, compute_loss_func=compute_loss, ) trainer.train() ``` ## Inference after training ```python item = val_ds[1] img = item["image"] bbox = json.loads(item["prompt"])["bbox"] inputs = processor(images=img, input_boxes=[[bbox]], return_tensors="pt").to( model.device ) with torch.no_grad(): outputs = model(**inputs) preds = processor.post_process_masks( outputs.pred_masks.cpu(), inputs["original_sizes"] )[0] preds = preds.squeeze(0) mask = (preds[0] > 0).cpu().numpy() overlay = np.asarray(img, dtype=np.uint8).copy() overlay[mask] = 0.55 * overlay[mask] + 0.45 * np.array([0, 255, 0], dtype=np.float32) plt.imshow(overlay) plt.title("After training") plt.axis("off") plt.show() ``` ```` ### `references/hub_saving.md` ````markdown # Saving Vision Models to Hugging Face Hub ## Contents - Why Hub Push is Required - Required Configuration (TrainingArguments, job config) - Complete Example - What Gets Saved - Important: Save Image Processor - Checkpoint Saving - Model Card Configuration - Saving Label Mappings - Authentication Methods - Verification Checklist - Repository Setup (automatic/manual creation, naming) - Troubleshooting (401, 403, push failures, inference issues) - Manual Push After Training - Example: Full Production Setup - Inference Example --- **CRITICAL:** Training environments are ephemeral. ALL results are lost when a job completes unless pushed to the Hub. ## Why Hub Push is Required When running on Hugging Face Jobs: - Environment is temporary - All files deleted on job completion - No local disk persistence - Cannot access results after job ends **Without Hub push, training is completely wasted.** ## Required Configuration ### 1. Training Configuration In your TrainingArguments: ```python from transformers import TrainingArguments training_args = TrainingArguments( output_dir="my-object-detector", push_to_hub=True, # Enable Hub push hub_model_id="username/model-name", # Target repository ) ``` ### 2. Job Configuration When submitting the job: ```python hf_jobs("uv", { "script": training_script_content, # Pass the Python script content directly as a string "secrets": {"HF_TOKEN": "$HF_TOKEN"} # Provide authentication }) ``` **The `$HF_TOKEN` syntax references your actual Hugging Face token value.** ## Complete Example ```python # train_detector.py # /// script # dependencies = ["transformers", "torch", "torchvision", "datasets"] # /// from transformers import ( AutoImageProcessor, AutoModelForObjectDetection, TrainingArguments, Trainer ) from datasets import load_dataset import os import torch # Load dataset dataset = load_dataset("cppe-5", split="train") # Load model and processor model_name = "facebook/detr-resnet-50" image_processor = AutoImageProcessor.from_pretrained(model_name) model = AutoModelForObjectDetection.from_pretrained( model_name, num_labels=5, # Number of classes ignore_mismatched_sizes=True ) # Configure with Hub push training_args = TrainingArguments( output_dir="my-detector", num_train_epochs=10, per_device_train_batch_size=8, # ✅ CRITICAL: Hub push configuration push_to_hub=True, hub_model_id="myusername/cppe5-detector", # Optional: Push strategy hub_strategy="checkpoint", # Push checkpoints during training ) # ✅ CRITICAL: Authenticate with Hub BEFORE creating Trainer from huggingface_hub import login hf_token = os.environ.get("HF_TOKEN") or os.environ.get("hfjob") if hf_token: login(token=hf_token) training_args.hub_token = hf_token elif training_args.push_to_hub: raise ValueError("HF_TOKEN not found! Add secrets={'HF_TOKEN': '$HF_TOKEN'} to job config.") # Define collate function def collate_fn(batch): pixel_values = [item["pixel_values"] for item in batch] labels = [item["labels"] for item in batch] encoding = image_processor.pad(pixel_values, return_tensors="pt") return { "pixel_values": encoding["pixel_values"], "labels": labels } trainer = Trainer( model=model, args=training_args, train_dataset=dataset, data_collator=collate_fn, ) trainer.train() # ✅ Push final model and processor trainer.push_to_hub() image_processor.push_to_hub("myusername/cppe5-detector") print("✅ Model saved to: https://huggingface.co/myusername/cppe5-detector") ``` **Submit with authentication:** ```python hf_jobs("uv", { "script": training_script_content, # Pass script content as a string, NOT a filename "flavor": "a10g-large", "timeout": "4h", "secrets": {"HF_TOKEN": "$HF_TOKEN"} # ✅ Required! }) ``` ## What Gets Saved When `push_to_hub=True`: 1. **Model weights** - Final trained parameters 2. **Image processor** - Associated preprocessing configuration 3. **Configuration** - Model config (config.json) including: - Number of labels/classes - Architecture details (backbone, num_queries, etc.) - Label mappings (id2label, label2id) 4. **Training arguments** - Hyperparameters used 5. **Model card** - Auto-generated documentation 6. **Checkpoints** - If `save_strategy="steps"` enabled ## Important: Save Image Processor **Object detection models require the image processor to be saved separately:** ```python # After training completes trainer.push_to_hub() # ✅ Also push the image processor image_processor.push_to_hub( repo_id="username/model-name", commit_message="Upload image processor" ) ``` **Why this matters:** - Models need specific image preprocessing (resizing, normalization) - Image processor contains critical configuration - Without it, model cannot be used for inference ## Checkpoint Saving Save intermediate checkpoints during training: ```python TrainingArguments( output_dir="my-detector", push_to_hub=True, hub_model_id="username/my-detector", # Checkpoint configuration save_strategy="steps", save_steps=500, # Save every 500 steps save_total_limit=3, # Keep only last 3 checkpoints hub_strategy="checkpoint", # Push checkpoints to Hub ) ``` **Benefits:** - Resume training if job fails - Compare checkpoint performance - Use intermediate models - Track training progress **Checkpoints are pushed to:** `username/my-detector` (same repo) ## Model Card Configuration Add metadata for better discoverability: ```python # At the end of training script model.push_to_hub( "username/my-detector", commit_message="Upload trained object detection model", tags=["object-detection", "vision", "cppe-5"], model_card_kwargs={ "license": "apache-2.0", "dataset": "cppe-5", "metrics": ["map", "recall", "precision"], "pipeline_tag": "object-detection", } ) ``` ## Saving Label Mappings **Critical for object detection:** Save class labels with the model: ```python # Define your label mappings id2label = {0: "Coverall", 1: "Face_Shield", 2: "Gloves", 3: "Goggles", 4: "Mask"} label2id = {v: k for k, v in id2label.items()} # Update model config before training model.config.id2label = id2label model.config.label2id = label2id # Now train and push trainer.train() trainer.push_to_hub() ``` **Without label mappings:** - Model outputs will be numeric IDs only - No human-readable class names - Difficult to interpret results ## Authentication Methods For a complete guide on token types, `$HF_TOKEN` automatic replacement, `secrets` vs `env` differences, and security best practices, see the `hugging-face-jobs` skill → *Token Usage Guide*. **Recommended:** Always pass tokens via `secrets` (encrypted server-side): ```python "secrets": {"HF_TOKEN": "$HF_TOKEN"} # ✅ Automatic replacement with your logged-in token ``` ## Verification Checklist Before submitting any training job, verify: - [ ] `push_to_hub=True` in TrainingArguments - [ ] `hub_model_id` is specified (format: `username/model-name`) - [ ] Image processor will be saved separately - [ ] Label mappings (id2label, label2id) are configured - [ ] Repository name doesn't conflict with existing repos - [ ] You have write access to the target namespace ## Repository Setup ### Automatic Creation If repository doesn't exist, it's created automatically when first pushing. ### Manual Creation Create repository before training: ```python from huggingface_hub import HfApi api = HfApi() api.create_repo( repo_id="username/detector-name", repo_type="model", private=False, # or True for private repo ) ``` ### Repository Naming **Valid names:** - `username/detr-cppe5` - `username/yolos-object-detector` - `organization/custom-detector` **Invalid names:** - `detector-name` (missing username) - `username/detector name` (spaces not allowed) - `username/DETECTOR` (uppercase discouraged) **Recommended naming:** - Include model architecture: `detr-`, `yolos-`, `deta-` - Include dataset: `-cppe5`, `-coco`, `-voc` - Be descriptive: `detr-resnet50-cppe5` > `model1` ## Troubleshooting ### Error: 401 Unauthorized **Cause:** HF_TOKEN not provided, invalid, or not authenticated before Trainer init **Solutions:** 1. Verify `secrets={"HF_TOKEN": "$HF_TOKEN"}` in job config 2. Verify script calls `login(token=hf_token)` AND sets `training_args.hub_token = hf_token` BEFORE creating the `Trainer` 3. Check you're logged in locally: `hf auth whoami` 4. Re-login: `hf auth login` **Root cause:** The `Trainer` calls `create_repo(token=self.args.hub_token)` during `__init__()` when `push_to_hub=True`. Relying on implicit env-var token resolution is unreliable in Jobs. Calling `login()` saves the token globally, and setting `training_args.hub_token` ensures the Trainer passes it explicitly to all Hub API calls. ### Error: 403 Forbidden **Cause:** No write access to repository **Solutions:** 1. Check repository namespace matches your username 2. Verify you're a member of organization (if using org namespace) 3. Check repository isn't private (if accessing org repo) ### Error: Repository not found **Cause:** Repository doesn't exist and auto-creation failed **Solutions:** 1. Manually create repository first 2. Check repository name format 3. Verify namespace exists ### Error: Push failed during training **Cause:** Network issues or Hub unavailable **Solutions:** 1. Training continues but final push fails 2. Checkpoints may be saved 3. Re-run push manually after job completes ### Issue: Model loads but inference fails **Possible causes:** 1. Image processor not saved—verify it's pushed separately 2. Label mappings missing—check config.json has id2label 3. Wrong image size—verify image processor matches training config ### Issue: Model saved but not visible **Possible causes:** 1. Repository is private—check https://huggingface.co/username 2. Wrong namespace—verify `hub_model_id` matches login 3. Push still in progress—wait a few minutes ## Manual Push After Training If training completes but push fails, push manually: ```python from transformers import AutoModelForObjectDetection, AutoImageProcessor # Load from local checkpoint model = AutoModelForObjectDetection.from_pretrained("./output_dir") image_processor = AutoImageProcessor.from_pretrained("./output_dir") # Push to Hub model.push_to_hub("username/model-name", token="hf_abc123...") image_processor.push_to_hub("username/model-name", token="hf_abc123...") ``` **Note:** Only possible if job hasn't completed (files still exist). ## Best Practices 1. **Always enable `push_to_hub=True`** 2. **Save image processor separately** - critical for inference 3. **Configure label mappings** before training 4. **Use checkpoint saving** for long training runs 5. **Verify Hub push** in logs before job completes 6. **Set appropriate `save_total_limit`** to avoid excessive checkpoints 7. **Use descriptive repo names** (e.g., `detr-cppe5` not `detector1`) 8. **Add model card** with: - Training dataset - Evaluation metrics (mAP, IoU) - Example usage code - Limitations 9. **Tag models appropriately**: - `object-detection` - Architecture: `detr`, `yolos`, `deta` - Dataset: `coco`, `voc`, `cppe-5` ## Monitoring Push Progress Check logs for push progress: ```python hf_jobs("logs", {"job_id": "your-job-id"}) ``` **Look for:** ``` Pushing model to username/detector-name... Upload file pytorch_model.bin: 100% ✅ Model pushed successfully Pushing image processor... ✅ Image processor pushed successfully ``` ## Example: Full Production Setup ```python # production_detector.py # /// script # dependencies = [ # "transformers>=4.30.0", # "torch>=2.0.0", # "torchvision>=0.15.0", # "datasets>=2.12.0", # "evaluate>=0.4.0" # ] # /// from transformers import ( AutoImageProcessor, AutoModelForObjectDetection, TrainingArguments, Trainer ) from datasets import load_dataset import os import torch # Configuration MODEL_NAME = "facebook/detr-resnet-50" DATASET_NAME = "cppe-5" HUB_MODEL_ID = "myusername/detr-cppe5-detector" NUM_CLASSES = 5 # Class labels id2label = {0: "Coverall", 1: "Face_Shield", 2: "Gloves", 3: "Goggles", 4: "Mask"} label2id = {v: k for k, v in id2label.items()} print(f"🔧 Loading dataset: {DATASET_NAME}") dataset = load_dataset(DATASET_NAME, split="train") print(f"✅ Dataset loaded: {len(dataset)} examples") print(f"🔧 Loading model: {MODEL_NAME}") image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME) model = AutoModelForObjectDetection.from_pretrained( MODEL_NAME, num_labels=NUM_CLASSES, id2label=id2label, label2id=label2id, ignore_mismatched_sizes=True ) print("✅ Model loaded") # Configure with comprehensive Hub settings training_args = TrainingArguments( output_dir="detr-cppe5", # Hub configuration push_to_hub=True, hub_model_id=HUB_MODEL_ID, hub_strategy="checkpoint", # Push checkpoints # Checkpoint configuration save_strategy="steps", save_steps=500, save_total_limit=3, # Training settings num_train_epochs=10, per_device_train_batch_size=8, gradient_accumulation_steps=2, learning_rate=1e-4, warmup_steps=500, # Evaluation eval_strategy="steps", eval_steps=500, # Logging logging_steps=50, logging_first_step=True, # Performance fp16=True, # Mixed precision training dataloader_num_workers=4, ) # ✅ CRITICAL: Authenticate with Hub BEFORE creating Trainer # login() saves the token globally so ALL hub operations can find it. from huggingface_hub import login hf_token = os.environ.get("HF_TOKEN") or os.environ.get("hfjob") if hf_token: login(token=hf_token) training_args.hub_token = hf_token elif training_args.push_to_hub: raise ValueError("HF_TOKEN not found! Add secrets={'HF_TOKEN': '$HF_TOKEN'} to job config.") # Data collator def collate_fn(batch): pixel_values = [item["pixel_values"] for item in batch] labels = [item["labels"] for item in batch] encoding = image_processor.pad(pixel_values, return_tensors="pt") return { "pixel_values": encoding["pixel_values"], "labels": labels } # Create trainer trainer = Trainer( model=model, args=training_args, train_dataset=dataset, data_collator=collate_fn, ) print("🚀 Starting training...") trainer.train() print("💾 Pushing final model to Hub...") trainer.push_to_hub( commit_message="Upload trained DETR model on CPPE-5", tags=["object-detection", "detr", "cppe-5", "vision"], ) print("💾 Pushing image processor to Hub...") image_processor.push_to_hub( repo_id=HUB_MODEL_ID, commit_message="Upload image processor" ) print("✅ Training complete!") print(f"Model available at: https://huggingface.co/{HUB_MODEL_ID}") print(f"\nTo use your model:") print(f"```python") print(f"from transformers import AutoImageProcessor, AutoModelForObjectDetection") print(f"") print(f"processor = AutoImageProcessor.from_pretrained('{HUB_MODEL_ID}')") print(f"model = AutoModelForObjectDetection.from_pretrained('{HUB_MODEL_ID}')") print(f"```") ``` **Submit:** ```python hf_jobs("uv", { "script": training_script_content, # Pass script content as a string, NOT a filename "flavor": "a10g-large", "timeout": "8h", "secrets": {"HF_TOKEN": "$HF_TOKEN"} }) ``` ## Inference Example After training, use your model: ```python from transformers import AutoImageProcessor, AutoModelForObjectDetection from PIL import Image import torch # Load model from Hub processor = AutoImageProcessor.from_pretrained("username/detr-cppe5-detector") model = AutoModelForObjectDetection.from_pretrained("username/detr-cppe5-detector") # Load and process image image = Image.open("test_image.jpg") inputs = processor(images=image, return_tensors="pt") # Run inference with torch.no_grad(): outputs = model(**inputs) # Post-process results target_sizes = torch.tensor([image.size[::-1]]) results = processor.post_process_object_detection( outputs, threshold=0.5, target_sizes=target_sizes )[0] # Print detections for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): box = [round(i, 2) for i in box.tolist()] print( f"Detected {model.config.id2label[label.item()]} with confidence " f"{round(score.item(), 3)} at location {box}" ) ``` ## Key Takeaway **Without `push_to_hub=True` and `secrets={"HF_TOKEN": "$HF_TOKEN"}`, all training results are permanently lost.** **For object detection, also remember to:** 1. Save the image processor separately 2. Configure label mappings (id2label, label2id) 3. Include appropriate model card metadata Always verify all three are configured before submitting any training job. ```` ### `references/image_classification_training_notebook.md` ````markdown # Image classification ## Contents - Load Food-101 dataset - Preprocess (ViT image processor, torchvision transforms) - Evaluate (accuracy metric, compute_metrics) - Train (TrainingArguments, Trainer setup, push to Hub) - Inference (pipeline, manual prediction) --- Image classification assigns a label or class to an image. Unlike text or audio classification, the inputs are the pixel values that comprise an image. There are many applications for image classification, such as detecting damage after a natural disaster, monitoring crop health, or helping screen medical images for signs of disease. This guide illustrates how to: 1. Fine-tune [ViT](../model_doc/vit) on the [Food-101](https://huggingface.co/datasets/ethz/food101) dataset to classify a food item in an image. 2. Use your fine-tuned model for inference. To see all architectures and checkpoints compatible with this task, we recommend checking the [task-page](https://huggingface.co/tasks/image-classification) Before you begin, make sure you have all the necessary libraries installed: ```bash pip install transformers datasets evaluate accelerate pillow torchvision scikit-learn trackio ``` We encourage you to log in to your Hugging Face account to upload and share your model with the community. When prompted, enter your token to log in: ```py >>> from huggingface_hub import notebook_login >>> notebook_login() ``` ## Load Food-101 dataset Start by loading a smaller subset of the Food-101 dataset from the 🤗 Datasets library. This will give you a chance to experiment and make sure everything works before spending more time training on the full dataset. ```py >>> from datasets import load_dataset >>> food = load_dataset("ethz/food101", split="train[:5000]") ``` Split the dataset's `train` split into a train and test set with the [train_test_split](https://huggingface.co/docs/datasets/v4.5.0/en/package_reference/main_classes#datasets.Dataset.train_test_split) method: ```py >>> food = food.train_test_split(test_size=0.2) ``` Then take a look at an example: ```py >>> food["train"][0] {'image': , 'label': 79} ``` Each example in the dataset has two fields: - `image`: a PIL image of the food item - `label`: the label class of the food item To make it easier for the model to get the label name from the label id, create a dictionary that maps the label name to an integer and vice versa: ```py >>> labels = food["train"].features["label"].names >>> label2id, id2label = dict(), dict() >>> for i, label in enumerate(labels): ... label2id[label] = str(i) ... id2label[str(i)] = label ``` Now you can convert the label id to a label name: ```py >>> id2label[str(79)] 'prime_rib' ``` ## Preprocess The next step is to load a ViT image processor to process the image into a tensor: ```py >>> from transformers import AutoImageProcessor >>> checkpoint = "google/vit-base-patch16-224-in21k" >>> image_processor = AutoImageProcessor.from_pretrained(checkpoint) ``` Apply some image transformations to the images to make the model more robust against overfitting. Here you'll use torchvision's [`transforms`](https://pytorch.org/vision/stable/transforms.html) module, but you can also use any image library you like. Crop a random part of the image, resize it, and normalize it with the image mean and standard deviation: ```py >>> from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor >>> normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std) >>> size = ( ... image_processor.size["shortest_edge"] ... if "shortest_edge" in image_processor.size ... else (image_processor.size["height"], image_processor.size["width"]) ... ) >>> _transforms = Compose([RandomResizedCrop(size), ToTensor(), normalize]) ``` Then create a preprocessing function to apply the transforms and return the `pixel_values` - the inputs to the model - of the image: ```py >>> def transforms(examples): ... examples["pixel_values"] = [_transforms(img.convert("RGB")) for img in examples["image"]] ... del examples["image"] ... return examples ``` To apply the preprocessing function over the entire dataset, use 🤗 Datasets [with_transform](https://huggingface.co/docs/datasets/v4.5.0/en/package_reference/main_classes#datasets.Dataset.with_transform) method. The transforms are applied on the fly when you load an element of the dataset: ```py >>> food = food.with_transform(transforms) ``` Now create a batch of examples using [DefaultDataCollator](/docs/transformers/v5.2.0/en/main_classes/data_collator#transformers.DefaultDataCollator). Unlike other data collators in 🤗 Transformers, the `DefaultDataCollator` does not apply additional preprocessing such as padding. ```py >>> from transformers import DefaultDataCollator >>> data_collator = DefaultDataCollator() ``` ## Evaluate Including a metric during training is often helpful for evaluating your model's performance. You can quickly load an evaluation method with the 🤗 [Evaluate](https://huggingface.co/docs/evaluate/index) library. For this task, load the [accuracy](https://huggingface.co/spaces/evaluate-metric/accuracy) metric (see the 🤗 Evaluate [quick tour](https://huggingface.co/docs/evaluate/a_quick_tour) to learn more about how to load and compute a metric): ```py >>> import evaluate >>> accuracy = evaluate.load("accuracy") ``` Then create a function that passes your predictions and labels to [compute](https://huggingface.co/docs/evaluate/v0.4.6/en/package_reference/main_classes#evaluate.EvaluationModule.compute) to calculate the accuracy: ```py >>> import numpy as np >>> def compute_metrics(eval_pred): ... predictions, labels = eval_pred ... predictions = np.argmax(predictions, axis=1) ... return accuracy.compute(predictions=predictions, references=labels) ``` Your `compute_metrics` function is ready to go now, and you'll return to it when you set up your training. ## Train If you aren't familiar with finetuning a model with the [Trainer](/docs/transformers/v5.2.0/en/main_classes/trainer#transformers.Trainer), take a look at the basic tutorial [here](../training#train-with-pytorch-trainer)! You're ready to start training your model now! Load ViT with [AutoModelForImageClassification](/docs/transformers/v5.2.0/en/model_doc/auto#transformers.AutoModelForImageClassification). Specify the number of labels along with the number of expected labels, and the label mappings: ```py >>> from transformers import AutoModelForImageClassification, TrainingArguments, Trainer >>> model = AutoModelForImageClassification.from_pretrained( ... checkpoint, ... num_labels=len(labels), ... id2label=id2label, ... label2id=label2id, ... ) ``` At this point, only three steps remain: 1. Define your training hyperparameters in [TrainingArguments](/docs/transformers/v5.2.0/en/main_classes/trainer#transformers.TrainingArguments). It is important you don't remove unused columns because that'll drop the `image` column. Without the `image` column, you can't create `pixel_values`. Set `remove_unused_columns=False` to prevent this behavior! The only other required parameter is `output_dir` which specifies where to save your model. You'll push this model to the Hub by setting `push_to_hub=True` (you need to be signed in to Hugging Face to upload your model). At the end of each epoch, the [Trainer](/docs/transformers/v5.2.0/en/main_classes/trainer#transformers.Trainer) will evaluate the accuracy and save the training checkpoint. 2. Pass the training arguments to [Trainer](/docs/transformers/v5.2.0/en/main_classes/trainer#transformers.Trainer) along with the model, dataset, tokenizer, data collator, and `compute_metrics` function. 3. Call [train()](/docs/transformers/v5.2.0/en/main_classes/trainer#transformers.Trainer.train) to finetune your model. ```py >>> training_args = TrainingArguments( ... output_dir="my_awesome_food_model", ... remove_unused_columns=False, ... eval_strategy="epoch", ... save_strategy="epoch", ... learning_rate=5e-5, ... per_device_train_batch_size=16, ... gradient_accumulation_steps=4, ... per_device_eval_batch_size=16, ... num_train_epochs=3, ... warmup_steps=0.1, ... logging_steps=10, ... report_to="trackio", ... run_name="food101", ... load_best_model_at_end=True, ... metric_for_best_model="accuracy", ... push_to_hub=True, ... ) >>> trainer = Trainer( ... model=model, ... args=training_args, ... data_collator=data_collator, ... train_dataset=food["train"], ... eval_dataset=food["test"], ... processing_class=image_processor, ... compute_metrics=compute_metrics, ... ) >>> trainer.train() ``` Once training is completed, share your model to the Hub with the [push_to_hub()](/docs/transformers/v5.2.0/en/main_classes/trainer#transformers.Trainer.push_to_hub) method so everyone can use your model: ```py >>> trainer.push_to_hub() ``` For a more in-depth example of how to finetune a model for image classification, take a look at the corresponding [PyTorch notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/image_classification.ipynb). ## Inference Great, now that you've fine-tuned a model, you can use it for inference! Load an image you'd like to run inference on: ```py >>> ds = load_dataset("ethz/food101", split="validation[:10]") >>> image = ds["image"][0] ``` The simplest way to try out your finetuned model for inference is to use it in a [pipeline()](/docs/transformers/v5.2.0/en/main_classes/pipelines#transformers.pipeline). Instantiate a `pipeline` for image classification with your model, and pass your image to it: ```py >>> from transformers import pipeline >>> classifier = pipeline("image-classification", model="my_awesome_food_model") >>> classifier(image) [{'score': 0.31856709718704224, 'label': 'beignets'}, {'score': 0.015232225880026817, 'label': 'bruschetta'}, {'score': 0.01519392803311348, 'label': 'chicken_wings'}, {'score': 0.013022331520915031, 'label': 'pork_chop'}, {'score': 0.012728818692266941, 'label': 'prime_rib'}] ``` You can also manually replicate the results of the `pipeline` if you'd like: Load an image processor to preprocess the image and return the `input` as PyTorch tensors: ```py >>> from transformers import AutoImageProcessor >>> import torch >>> image_processor = AutoImageProcessor.from_pretrained("my_awesome_food_model") >>> inputs = image_processor(image, return_tensors="pt") ``` Pass your inputs to the model and return the logits: ```py >>> from transformers import AutoModelForImageClassification >>> model = AutoModelForImageClassification.from_pretrained("my_awesome_food_model") >>> with torch.no_grad(): ... logits = model(**inputs).logits ``` Get the predicted label with the highest probability, and use the model's `id2label` mapping to convert it to a label: ```py >>> predicted_label = logits.argmax(-1).item() >>> model.config.id2label[predicted_label] 'beignets' ``` ```` ### `references/object_detection_training_notebook.md` ````markdown # Object Detection Training Reference ## Contents - Load the CPPE-5 dataset - Preprocess the data (augmentation with Albumentations, COCO annotation formatting) - Preparing function to compute mAP - Training the detection model (TrainingArguments, Trainer setup) - Evaluate - Inference (loading from Hub, running predictions, visualizing results) --- Object detection is the computer vision task of detecting instances (such as humans, buildings, or cars) in an image. Object detection models receive an image as input and output coordinates of the bounding boxes and associated labels of the detected objects. An image can contain multiple objects, each with its own bounding box and a label (e.g. it can have a car and a building), and each object can be present in different parts of an image (e.g. the image can have several cars). This task is commonly used in autonomous driving for detecting things like pedestrians, road signs, and traffic lights. Other applications include counting objects in images, image search, and more. In this guide, you will learn how to: 1. Finetune [DETR](https://huggingface.co/docs/transformers/model_doc/detr), a model that combines a convolutional backbone with an encoder-decoder Transformer, on the [CPPE-5](https://huggingface.co/datasets/cppe-5) dataset. 2. Use your finetuned model for inference. To see all architectures and checkpoints compatible with this task, we recommend checking the [task-page](https://huggingface.co/tasks/object-detection) Before you begin, make sure you have all the necessary libraries installed: ```bash pip install -q datasets transformers accelerate timm trackio pip install -q -U albumentations>=1.4.5 torchmetrics pycocotools ``` You'll use 🤗 Datasets to load a dataset from the Hugging Face Hub, 🤗 Transformers to train your model, and `albumentations` to augment the data. We encourage you to share your model with the community. Log in to your Hugging Face account to upload it to the Hub. When prompted, enter your token to log in: ```py >>> from huggingface_hub import notebook_login >>> notebook_login() ``` To get started, we'll define global constants, namely the model name and image size. For this tutorial, we'll use the conditional DETR model due to its faster convergence. Feel free to select any object detection model available in the `transformers` library. ```py >>> MODEL_NAME = "microsoft/conditional-detr-resnet-50" # or "facebook/detr-resnet-50" >>> IMAGE_SIZE = 480 ``` ## Load the CPPE-5 dataset The [CPPE-5 dataset](https://huggingface.co/datasets/cppe-5) contains images with annotations identifying medical personal protective equipment (PPE) in the context of the COVID-19 pandemic. Start by loading the dataset and creating a `validation` split from `train`: ```py >>> from datasets import load_dataset >>> cppe5 = load_dataset("cppe-5") >>> if "validation" not in cppe5: ... split = cppe5["train"].train_test_split(0.15, seed=1337) ... cppe5["train"] = split["train"] ... cppe5["validation"] = split["test"] >>> cppe5 DatasetDict({ train: Dataset({ features: ['image_id', 'image', 'width', 'height', 'objects'], num_rows: 850 }) test: Dataset({ features: ['image_id', 'image', 'width', 'height', 'objects'], num_rows: 29 }) validation: Dataset({ features: ['image_id', 'image', 'width', 'height', 'objects'], num_rows: 150 }) }) ``` You'll see that this dataset has 1000 images for train and validation sets and a test set with 29 images. To get familiar with the data, explore what the examples look like. ```py >>> cppe5["train"][0] { 'image_id': 366, 'image': , 'width': 500, 'height': 500, 'objects': { 'id': [1932, 1933, 1934], 'area': [27063, 34200, 32431], 'bbox': [[29.0, 11.0, 97.0, 279.0], [201.0, 1.0, 120.0, 285.0], [382.0, 0.0, 113.0, 287.0]], 'category': [0, 0, 0] } } ``` The examples in the dataset have the following fields: - `image_id`: the example image id - `image`: a `PIL.Image.Image` object containing the image - `width`: width of the image - `height`: height of the image - `objects`: a dictionary containing bounding box metadata for the objects in the image: - `id`: the annotation id - `area`: the area of the bounding box - `bbox`: the object's bounding box (in the [COCO format](https://albumentations.ai/docs/getting_started/bounding_boxes_augmentation/#coco) ) - `category`: the object's category, with possible values including `Coverall (0)`, `Face_Shield (1)`, `Gloves (2)`, `Goggles (3)` and `Mask (4)` You may notice that the `bbox` field follows the COCO format, which is the format that the DETR model expects. However, the grouping of the fields inside `objects` differs from the annotation format DETR requires. You will need to apply some preprocessing transformations before using this data for training. To get an even better understanding of the data, visualize an example in the dataset. ```py >>> import numpy as np >>> import os >>> from PIL import Image, ImageDraw >>> image = cppe5["train"][2]["image"] >>> annotations = cppe5["train"][2]["objects"] >>> draw = ImageDraw.Draw(image) >>> categories = cppe5["train"].features["objects"]["category"].feature.names >>> id2label = {index: x for index, x in enumerate(categories, start=0)} >>> label2id = {v: k for k, v in id2label.items()} >>> for i in range(len(annotations["id"])): ... box = annotations["bbox"][i] ... class_idx = annotations["category"][i] ... x, y, w, h = tuple(box) ... # Check if coordinates are normalized or not ... if max(box) > 1.0: ... # Coordinates are un-normalized, no need to re-scale them ... x1, y1 = int(x), int(y) ... x2, y2 = int(x + w), int(y + h) ... else: ... # Coordinates are normalized, re-scale them ... x1 = int(x * width) ... y1 = int(y * height) ... x2 = int((x + w) * width) ... y2 = int((y + h) * height) ... draw.rectangle((x, y, x + w, y + h), outline="red", width=1) ... draw.text((x, y), id2label[class_idx], fill="white") >>> image ``` To visualize the bounding boxes with associated labels, you can get the labels from the dataset's metadata, specifically the `category` field. You'll also want to create dictionaries that map a label id to a label class (`id2label`) and the other way around (`label2id`). You can use them later when setting up the model. Including these maps will make your model reusable by others if you share it on the Hugging Face Hub. Please note that, the part of above code that draws the bounding boxes assume that it is in `COCO` format `(x_min, y_min, width, height)`. It has to be adjusted to work for other formats like `(x_min, y_min, x_max, y_max)`. As a final step of getting familiar with the data, explore it for potential issues. One common problem with datasets for object detection is bounding boxes that "stretch" beyond the edge of the image. Such "runaway" bounding boxes can raise errors during training and should be addressed. There are a few examples with this issue in this dataset. To keep things simple in this guide, we will set `clip=True` for `BboxParams` in transformations below. ## Preprocess the data To finetune a model, you must preprocess the data you plan to use to match precisely the approach used for the pre-trained model. [AutoImageProcessor](/docs/transformers/v5.1.0/en/model_doc/auto#transformers.AutoImageProcessor) takes care of processing image data to create `pixel_values`, `pixel_mask`, and `labels` that a DETR model can train with. The image processor has some attributes that you won't have to worry about: - `image_mean = [0.485, 0.456, 0.406 ]` - `image_std = [0.229, 0.224, 0.225]` These are the mean and standard deviation used to normalize images during the model pre-training. These values are crucial to replicate when doing inference or finetuning a pre-trained image model. Instantiate the image processor from the same checkpoint as the model you want to finetune. ```py >>> from transformers import AutoImageProcessor >>> MAX_SIZE = IMAGE_SIZE >>> image_processor = AutoImageProcessor.from_pretrained( ... MODEL_NAME, ... do_resize=True, ... size={"max_height": MAX_SIZE, "max_width": MAX_SIZE}, ... do_pad=True, ... pad_size={"height": MAX_SIZE, "width": MAX_SIZE}, ... ) ``` Before passing the images to the `image_processor`, apply two preprocessing transformations to the dataset: - Augmenting images - Reformatting annotations to meet DETR expectations First, to make sure the model does not overfit on the training data, you can apply image augmentation with any data augmentation library. Here we use [Albumentations](https://albumentations.ai/docs/). This library ensures that transformations affect the image and update the bounding boxes accordingly. The 🤗 Datasets library documentation has a detailed [guide on how to augment images for object detection](https://huggingface.co/docs/datasets/object_detection), and it uses the exact same dataset as an example. Apply some geometric and color transformations to the image. For additional augmentation options, explore the [Albumentations Demo Space](https://huggingface.co/spaces/qubvel-hf/albumentations-demo). ```py >>> import albumentations as A >>> train_augment_and_transform = A.Compose( ... [ ... A.Perspective(p=0.1), ... A.HorizontalFlip(p=0.5), ... A.RandomBrightnessContrast(p=0.5), ... A.HueSaturationValue(p=0.1), ... ], ... bbox_params=A.BboxParams(format="coco", label_fields=["category"], clip=True, min_area=25), ... ) >>> validation_transform = A.Compose( ... [A.NoOp()], ... bbox_params=A.BboxParams(format="coco", label_fields=["category"], clip=True), ... ) ``` The `image_processor` expects the annotations to be in the following format: `{'image_id': int, 'annotations': list[Dict]}`, where each dictionary is a COCO object annotation. Let's add a function to reformat annotations for a single example: ```py >>> def format_image_annotations_as_coco(image_id, categories, areas, bboxes): ... """Format one set of image annotations to the COCO format ... Args: ... image_id (str): image id. e.g. "0001" ... categories (list[int]): list of categories/class labels corresponding to provided bounding boxes ... areas (list[float]): list of corresponding areas to provided bounding boxes ... bboxes (list[tuple[float]]): list of bounding boxes provided in COCO format ... ([center_x, center_y, width, height] in absolute coordinates) ... Returns: ... dict: { ... "image_id": image id, ... "annotations": list of formatted annotations ... } ... """ ... annotations = [] ... for category, area, bbox in zip(categories, areas, bboxes): ... formatted_annotation = { ... "image_id": image_id, ... "category_id": category, ... "iscrowd": 0, ... "area": area, ... "bbox": list(bbox), ... } ... annotations.append(formatted_annotation) ... return { ... "image_id": image_id, ... "annotations": annotations, ... } ``` Now you can combine the image and annotation transformations to use on a batch of examples: ```py >>> def augment_and_transform_batch(examples, transform, image_processor, return_pixel_mask=False): ... """Apply augmentations and format annotations in COCO format for object detection task""" ... images = [] ... annotations = [] ... for image_id, image, objects in zip(examples["image_id"], examples["image"], examples["objects"]): ... image = np.array(image.convert("RGB")) ... # apply augmentations ... output = transform(image=image, bboxes=objects["bbox"], category=objects["category"]) ... images.append(output["image"]) ... # format annotations in COCO format ... formatted_annotations = format_image_annotations_as_coco( ... image_id, output["category"], objects["area"], output["bboxes"] ... ) ... annotations.append(formatted_annotations) ... # Apply the image processor transformations: resizing, rescaling, normalization ... result = image_processor(images=images, annotations=annotations, return_tensors="pt") ... if not return_pixel_mask: ... result.pop("pixel_mask", None) ... return result ``` Apply this preprocessing function to the entire dataset using 🤗 Datasets [with_transform](https://huggingface.co/docs/datasets/v4.5.0/en/package_reference/main_classes#datasets.Dataset.with_transform) method. This method applies transformations on the fly when you load an element of the dataset. At this point, you can check what an example from the dataset looks like after the transformations. You should see a tensor with `pixel_values`, a tensor with `pixel_mask`, and `labels`. ```py >>> from functools import partial >>> # Make transform functions for batch and apply for dataset splits >>> train_transform_batch = partial( ... augment_and_transform_batch, transform=train_augment_and_transform, image_processor=image_processor ... ) >>> validation_transform_batch = partial( ... augment_and_transform_batch, transform=validation_transform, image_processor=image_processor ... ) >>> cppe5["train"] = cppe5["train"].with_transform(train_transform_batch) >>> cppe5["validation"] = cppe5["validation"].with_transform(validation_transform_batch) >>> cppe5["test"] = cppe5["test"].with_transform(validation_transform_batch) >>> cppe5["train"][15] {'pixel_values': tensor([[[ 1.9235, 1.9407, 1.9749, ..., -0.7822, -0.7479, -0.6965], [ 1.9578, 1.9749, 1.9920, ..., -0.7993, -0.7650, -0.7308], [ 2.0092, 2.0092, 2.0263, ..., -0.8507, -0.8164, -0.7822], ..., [ 0.0741, 0.0741, 0.0741, ..., 0.0741, 0.0741, 0.0741], [ 0.0741, 0.0741, 0.0741, ..., 0.0741, 0.0741, 0.0741], [ 0.0741, 0.0741, 0.0741, ..., 0.0741, 0.0741, 0.0741]], [[ 1.6232, 1.6408, 1.6583, ..., 0.8704, 1.0105, 1.1331], [ 1.6408, 1.6583, 1.6758, ..., 0.8529, 0.9930, 1.0980], [ 1.6933, 1.6933, 1.7108, ..., 0.8179, 0.9580, 1.0630], ..., [ 0.2052, 0.2052, 0.2052, ..., 0.2052, 0.2052, 0.2052], [ 0.2052, 0.2052, 0.2052, ..., 0.2052, 0.2052, 0.2052], [ 0.2052, 0.2052, 0.2052, ..., 0.2052, 0.2052, 0.2052]], [[ 1.8905, 1.9080, 1.9428, ..., -0.1487, -0.0964, -0.0615], [ 1.9254, 1.9428, 1.9603, ..., -0.1661, -0.1138, -0.0790], [ 1.9777, 1.9777, 1.9951, ..., -0.2010, -0.1138, -0.0790], ..., [ 0.4265, 0.4265, 0.4265, ..., 0.4265, 0.4265, 0.4265], [ 0.4265, 0.4265, 0.4265, ..., 0.4265, 0.4265, 0.4265], [ 0.4265, 0.4265, 0.4265, ..., 0.4265, 0.4265, 0.4265]]]), 'labels': {'image_id': tensor([688]), 'class_labels': tensor([3, 4, 2, 0, 0]), 'boxes': tensor([[0.4700, 0.1933, 0.1467, 0.0767], [0.4858, 0.2600, 0.1150, 0.1000], [0.4042, 0.4517, 0.1217, 0.1300], [0.4242, 0.3217, 0.3617, 0.5567], [0.6617, 0.4033, 0.5400, 0.4533]]), 'area': tensor([ 4048., 4140., 5694., 72478., 88128.]), 'iscrowd': tensor([0, 0, 0, 0, 0]), 'orig_size': tensor([480, 480])}} ``` You have successfully augmented the individual images and prepared their annotations. However, preprocessing isn't complete yet. In the final step, create a custom `collate_fn` to batch images together. Pad images (which are now `pixel_values`) to the largest image in a batch, and create a corresponding `pixel_mask` to indicate which pixels are real (1) and which are padding (0). ```py >>> import torch >>> def collate_fn(batch): ... data = {} ... data["pixel_values"] = torch.stack([x["pixel_values"] for x in batch]) ... data["labels"] = [x["labels"] for x in batch] ... if "pixel_mask" in batch[0]: ... data["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch]) ... return data ``` ## Preparing function to compute mAP Object detection models are commonly evaluated with a set of COCO-style metrics. We are going to use `torchmetrics` to compute `mAP` (mean average precision) and `mAR` (mean average recall) metrics and will wrap it to `compute_metrics` function in order to use in [Trainer](/docs/transformers/v5.1.0/en/main_classes/trainer#transformers.Trainer) for evaluation. Intermediate format of boxes used for training is `YOLO` (normalized) but we will compute metrics for boxes in `Pascal VOC` (absolute) format in order to correctly handle box areas. Let's define a function that converts bounding boxes to `Pascal VOC` format: ```py >>> from transformers.image_transforms import center_to_corners_format >>> def convert_bbox_yolo_to_pascal(boxes, image_size): ... """ ... Convert bounding boxes from YOLO format (x_center, y_center, width, height) in range [0, 1] ... to Pascal VOC format (x_min, y_min, x_max, y_max) in absolute coordinates. ... Args: ... boxes (torch.Tensor): Bounding boxes in YOLO format ... image_size (tuple[int, int]): Image size in format (height, width) ... Returns: ... torch.Tensor: Bounding boxes in Pascal VOC format (x_min, y_min, x_max, y_max) ... """ ... # convert center to corners format ... boxes = center_to_corners_format(boxes) ... # convert to absolute coordinates ... height, width = image_size ... boxes = boxes * torch.tensor([[width, height, width, height]]) ... return boxes ``` Then, in `compute_metrics` function we collect `predicted` and `target` bounding boxes, scores and labels from evaluation loop results and pass it to the scoring function. ```py >>> import numpy as np >>> from dataclasses import dataclass >>> from torchmetrics.detection.mean_ap import MeanAveragePrecision >>> @dataclass >>> class ModelOutput: ... logits: torch.Tensor ... pred_boxes: torch.Tensor >>> @torch.no_grad() >>> def compute_metrics(evaluation_results, image_processor, threshold=0.0, id2label=None): ... """ ... Compute mean average mAP, mAR and their variants for the object detection task. ... Args: ... evaluation_results (EvalPrediction): Predictions and targets from evaluation. ... threshold (float, optional): Threshold to filter predicted boxes by confidence. Defaults to 0.0. ... id2label (Optional[dict], optional): Mapping from class id to class name. Defaults to None. ... Returns: ... Mapping[str, float]: Metrics in a form of dictionary {: } ... """ ... predictions, targets = evaluation_results.predictions, evaluation_results.label_ids ... # For metric computation we need to provide: ... # - targets in a form of list of dictionaries with keys "boxes", "labels" ... # - predictions in a form of list of dictionaries with keys "boxes", "scores", "labels" ... image_sizes = [] ... post_processed_targets = [] ... post_processed_predictions = [] ... # Collect targets in the required format for metric computation ... for batch in targets: ... # collect image sizes, we will need them for predictions post processing ... batch_image_sizes = torch.tensor(np.array([x["orig_size"] for x in batch])) ... image_sizes.append(batch_image_sizes) ... # collect targets in the required format for metric computation ... # boxes were converted to YOLO format needed for model training ... # here we will convert them to Pascal VOC format (x_min, y_min, x_max, y_max) ... for image_target in batch: ... boxes = torch.tensor(image_target["boxes"]) ... boxes = convert_bbox_yolo_to_pascal(boxes, imag ... ```` ### `references/reliability_principles.md` ````markdown # Reliability Principles for Training Jobs ## Contents - Principle 1: Always Verify Before Use - Principle 2: Prioritize Reliability Over Performance - Principle 3: Create Atomic, Self-Contained Scripts - Principle 4: Provide Clear Error Context - Principle 5: Test the Happy Path on Known-Good Inputs - Summary: The Reliability Checklist (pre-flight, script quality, job config) - When Principles Conflict --- These principles are derived from real production failures and successful fixes. Following them prevents common failure modes and ensures reliable job execution. ## Principle 1: Always Verify Before Use **Rule:** Never assume repos, datasets, or resources exist. Verify with tools first. ### What It Prevents - **Non-existent datasets** - Jobs fail immediately when dataset doesn't exist - **Typos in names** - Simple mistakes like "argilla-dpo-mix-7k" vs "ultrafeedback_binarized" - **Incorrect paths** - Old or moved repos, renamed files - **Missing dependencies** - Undocumented requirements ### How to Apply **Before submitting ANY job:** ```python # Verify dataset exists dataset_search({"query": "dataset-name", "author": "author-name", "limit": 5}) hub_repo_details(["author/dataset-name"], repo_type="dataset") # Verify model exists hub_repo_details(["org/model-name"], repo_type="model") # Check script/file paths (for URL-based scripts) # Verify before using: https://github.com/user/repo/blob/main/script.py ``` **Examples that would have caught errors:** ```python # ❌ WRONG: Assumed dataset exists hf_jobs("uv", { "script": """...""", "env": {"DATASET": "trl-lib/argilla-dpo-mix-7k"} # Doesn't exist! }) # ✅ CORRECT: Verify first dataset_search({"query": "argilla dpo", "author": "trl-lib"}) # Would show: "trl-lib/ultrafeedback_binarized" is the correct name hub_repo_details(["trl-lib/ultrafeedback_binarized"], repo_type="dataset") # Confirms it exists before using ``` ### Implementation Checklist - [ ] Check dataset exists before training - [ ] Test script URLs are valid before submitting - [ ] Check for recent updates/renames of resources - [ ] Check for dataset format **Time cost:** 5-10 seconds **Time saved:** Hours of failed job time + debugging --- ## Principle 2: Prioritize Reliability Over Performance **Rule:** Default to what is most likely to succeed, not what is theoretically fastest. ### What It Prevents - **Hardware incompatibilities** - Features that fail on certain GPUs - **Unstable optimizations** - Speed-ups that cause crashes - **Complex configurations** - More failure points - **Build system issues** - Unreliable compilation methods ### How to Apply **Choose reliability:** ```python # ❌ RISKY: Aggressive optimization that may fail TrainingArguments( torch_compile=True, # Can fail on T4, A10G GPUs optim="adamw_bnb_8bit", # Requires specific setup dataloader_num_workers=8, # May cause OOM on small instances ... ) # ✅ SAFE: Proven defaults TrainingArguments( # torch_compile=True, # Commented with note: "Enable on H100 for 20% speedup" optim="adamw_torch", # Standard, always works fp16=True, # Stable and fast on T4/A10G dataloader_num_workers=4, # Conservative, reliable ... ) ``` ### Real-World Example **The `torch.compile` failure:** - Added for "20% speedup" on H100 - **Failed fatally on T4-medium** with cryptic error - Misdiagnosed as dataset issue (cost hours) - **Fix:** Disable by default, add as optional comment **Result:** Reliability > 20% performance gain ### Implementation Checklist - [ ] Use proven, standard configurations by default - [ ] Comment out performance optimizations with hardware notes - [ ] Use stable build systems (CMake > make) - [ ] Test on target hardware before production - [ ] Document known incompatibilities - [ ] Provide "safe" and "fast" variants when needed **Performance loss:** 10-20% in best case **Reliability gain:** 95%+ success rate vs 60-70% --- ## Principle 3: Create Atomic, Self-Contained Scripts **Rule:** Scripts should work as complete, independent units. Don't remove parts to "simplify." ### What It Prevents - **Missing dependencies** - Removed "unnecessary" packages that are actually required - **Incomplete processes** - Skipped steps that seem redundant - **Environment assumptions** - Scripts that need pre-setup - **Partial failures** - Some parts work, others fail silently ### How to Apply **Complete dependency specifications:** ```python # ❌ INCOMPLETE: "Simplified" by removing dependencies # /// script # dependencies = [ # "transformers", # "torch", # "datasets", # ] # /// # ✅ COMPLETE: All dependencies explicit # /// script # dependencies = [ # "transformers>=5.2.0", # "accelerate>=1.1.0", # "albumentations>=1.4.16", # Required for augmentation + bbox handling # "timm", # Required for vision backbones # "datasets>=4.0", # "torchmetrics", # Required for mAP/mAR computation # "pycocotools", # Required for COCO evaluation # "trackio", # Required for metrics monitoring # "huggingface_hub", # ] # /// ``` ### Real-World Example **The `albumentations` failure:** - Original script had it: augmentations and bbox clipping worked fine - "Simplified" version removed it: "not strictly needed for training" - **Training crashed on bbox augmentation** — no fallback for COCO-format bbox handling - Hard to debug: error appeared in data loading, not in augmentation setup - **Fix:** Restore all original dependencies **Result:** Don't remove dependencies without thorough testing ### Implementation Checklist - [ ] All dependencies in PEP 723 header with version pins - [ ] All system packages installed by script - [ ] No assumptions about pre-existing environment - [ ] No "optional" steps that are actually required - [ ] Test scripts in clean environment - [ ] Document why each dependency is needed **Complexity:** Slightly longer scripts **Reliability:** Scripts "just work" every time --- ## Principle 4: Provide Clear Error Context **Rule:** When things fail, make it obvious what went wrong and how to fix it. ### How to Apply **Wrap subprocess calls:** ```python # ❌ UNCLEAR: Silent failure subprocess.run([...], check=True, capture_output=True) # ✅ CLEAR: Shows what failed try: result = subprocess.run( [...], check=True, capture_output=True, text=True ) print(result.stdout) if result.stderr: print("Warnings:", result.stderr) except subprocess.CalledProcessError as e: print(f"❌ Command failed!") print("STDOUT:", e.stdout) print("STDERR:", e.stderr) raise ``` **Validate inputs:** ```python # ❌ UNCLEAR: Fails later with cryptic error model = load_model(MODEL_NAME) # ✅ CLEAR: Fails fast with clear message if not MODEL_NAME: raise ValueError("MODEL_NAME environment variable not set!") print(f"Loading model: {MODEL_NAME}") try: model = load_model(MODEL_NAME) print(f"✅ Model loaded successfully") except Exception as e: print(f"❌ Failed to load model: {MODEL_NAME}") print(f"Error: {e}") print("Hint: Check that model exists on Hub") raise ``` ### Implementation Checklist - [ ] Wrap external calls with try/except - [ ] Print stdout/stderr on failure - [ ] Validate environment variables early - [ ] Add progress indicators (✅, ❌, 🔄) - [ ] Include hints for common failures - [ ] Log configuration at start --- ## Principle 5: Test the Happy Path on Known-Good Inputs **Rule:** Before using new code in production, test with inputs you know work. ## Summary: The Reliability Checklist Before submitting ANY job: ### Pre-Flight Checks - [ ] **Verified** all repos/datasets exist (hub_repo_details) - [ ] **Tested** with known-good inputs if new code - [ ] **Using** proven hardware/configuration - [ ] **Included** all dependencies in PEP 723 header - [ ] **Installed** system requirements (build tools, etc.) - [ ] **Set** appropriate timeout (not default 30m) - [ ] **Configured** Hub push with HF_TOKEN (login() + hub_token) - [ ] **Added** clear error handling ### Script Quality - [ ] Self-contained (no external setup needed) - [ ] Complete dependencies listed - [ ] Build tools installed by script - [ ] Progress indicators included - [ ] Error messages are clear - [ ] Configuration logged at start ### Job Configuration - [ ] Timeout > expected runtime + 30% buffer - [ ] Hardware appropriate for model size - [ ] Secrets include HF_TOKEN (see SKILL.md directive #2 for syntax) - [ ] Script calls `login(token=hf_token)` and sets `training_args.hub_token = hf_token` BEFORE `Trainer()` init - [ ] Environment variables set correctly - [ ] Cost estimated and acceptable **Following these principles transforms job success rate from ~60-70% to ~95%+** --- ## When Principles Conflict Sometimes reliability and performance conflict. Here's how to choose: | Scenario | Choose | Rationale | |----------|--------|-----------| | Demo/test | Reliability | Fast failure is worse than slow success | | Production (first run) | Reliability | Prove it works before optimizing | | Production (proven) | Performance | Safe to optimize after validation | | Time-critical | Reliability | Failures cause more delay than slow runs | | Cost-critical | Balanced | Test with small model, then optimize | **General rule:** Reliability first, optimize second. --- ```` ### `references/timm_trainer.md` ````markdown # Using timm models with Hugging Face Trainer Transformers has first-class support for timm models via the `TimmWrapper` classes. You can load any timm model and use it directly with the `Trainer` API for image classification. Here's how it works: ## Loading a timm model The `TimmWrapperForImageClassification` class (in `transformers/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py`) wraps timm models so they're fully compatible with the Trainer API. You can load them via the `Auto` classes: ```python from transformers import AutoModelForImageClassification, AutoImageProcessor, Trainer, TrainingArguments # Load a timm model for image classification checkpoint = "timm/resnet50.a1_in1k" image_processor = AutoImageProcessor.from_pretrained(checkpoint) model = AutoModelForImageClassification.from_pretrained( checkpoint, num_labels=10, # set to your number of classes ignore_mismatched_sizes=True, # needed when changing num_labels from pretrained ) ``` ## Key details 1. **Image processor**: The `TimmWrapperImageProcessor` automatically resolves the correct transforms from timm's config. It exposes both `val_transforms` and `train_transforms` (with augmentations), as noted in the code: ```64:65:transformers/src/transformers/models/timm_wrapper/image_processing_timm_wrapper.py # useful for training, see examples/pytorch/image-classification/run_image_classification.py self.train_transforms = timm.data.create_transform(**self.data_config, is_training=True) ``` 2. **Loss computation is built-in**: `TimmWrapperForImageClassification.forward()` accepts a `labels` argument and computes cross-entropy loss automatically, which is exactly what Trainer expects: ```374:376:transformers/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py loss = None if labels is not None: loss = self.loss_function(labels, logits, self.config) ``` 3. **Returns `ImageClassifierOutput`**: The output format is the standard transformers output, so Trainer handles it seamlessly. ## Full training example ```python from transformers import AutoModelForImageClassification, AutoImageProcessor, Trainer, TrainingArguments from datasets import load_dataset # Load dataset dataset = load_dataset("food101", split="train[:5000]") dataset = dataset.train_test_split(test_size=0.2) # Load timm model + processor checkpoint = "timm/resnet50.a1_in1k" image_processor = AutoImageProcessor.from_pretrained(checkpoint) model = AutoModelForImageClassification.from_pretrained( checkpoint, num_labels=101, ignore_mismatched_sizes=True, ) # Preprocessing def transform(batch): batch["pixel_values"] = [image_processor(img)["pixel_values"][0] for img in batch["image"]] batch["labels"] = batch["label"] return batch dataset["train"].set_transform(transform) dataset["test"].set_transform(transform) # Train training_args = TrainingArguments( output_dir="./timm-finetuned", num_train_epochs=3, per_device_train_batch_size=16, per_device_eval_batch_size=16, eval_strategy="epoch", save_strategy="epoch", logging_steps=50, remove_unused_columns=False, ) trainer = Trainer( model=model, args=training_args, train_dataset=dataset["train"], eval_dataset=dataset["test"], ) trainer.train() ``` Any timm checkpoint on the Hub (prefixed with `timm/`) works out of the box (ResNet, EfficientNet, ViT, ConvNeXt, etc). The wrapper handles all the translation between timm's interface and what Trainer expects. ```` ### `scripts/dataset_inspector.py` ``` #!/usr/bin/env python3 # /// script # requires-python = ">=3.10" # dependencies = [] # /// """ Dataset Format Inspector for Vision Model Training Inspects Hugging Face datasets to determine compatibility with object detection and image classification training. Uses Datasets Server API for instant results - no dataset download needed! ULTRA-EFFICIENT: Uses HF Datasets Server API - completes in <2 seconds. Usage with HF Jobs: hf_jobs("uv", { "script": "path/to/dataset_inspector.py", "script_args": ["--dataset", "your/dataset", "--split", "train"] }) """ import argparse import math import sys import json import urllib.request import urllib.parse from typing import List, Dict, Any, Tuple def parse_args(): parser = argparse.ArgumentParser(description="Inspect dataset format for vision model training") parser.add_argument("--dataset", type=str, required=True, help="Dataset name") parser.add_argument("--split", type=str, default="train", help="Dataset split (default: train)") parser.add_argument("--config", type=str, default="default", help="Dataset config name (default: default)") parser.add_argument("--preview", type=int, default=150, help="Max chars per field preview") parser.add_argument("--samples", type=int, default=5, help="Number of samples to fetch (default: 5)") parser.add_argument("--json-output", action="store_true", help="Output as JSON") return parser.parse_args() def api_request(url: str) -> Dict: """Make API request to Datasets Server""" try: with urllib.request.urlopen(url, timeout=10) as response: return json.loads(response.read().decode()) except urllib.error.HTTPError as e: if e.code == 404: return None raise Exception(f"API request failed: {e.code} {e.reason}") except Exception as e: raise Exception(f"API request failed: {str(e)}") def get_splits(dataset: str) -> Dict: """Get available splits for dataset""" url = f"https://datasets-server.huggingface.co/splits?dataset={urllib.parse.quote(dataset)}" return api_request(url) def get_rows(dataset: str, config: str, split: str, offset: int = 0, length: int = 5) -> Dict: """Get rows from dataset""" url = f"https://datasets-server.huggingface.co/rows?dataset={urllib.parse.quote(dataset)}&config={config}&split={split}&offset={offset}&length={length}" return api_request(url) def find_columns(columns: List[str], patterns: List[str]) -> List[str]: """Find columns matching patterns""" return [c for c in columns if any(p in c.lower() for p in patterns)] def detect_bbox_format(bbox: List[float], image_size: Tuple[int, int] = None) -> str: """ Detect bounding box format based on values and optionally image dimensions. Common formats: - [x_min, y_min, x_max, y_max] - XYXY (Pascal VOC) - [x_min, y_min, width, height] - XYWH (COCO) - [x_center, y_center, width, height] - CXCYWH (YOLO normalized) """ if len(bbox) != 4: return "unknown (not 4 values)" a, b, c, d = bbox is_normalized = all(0 <= v <= 1 for v in bbox) if c < a or d < b: if is_normalized: return "xywh_normalized" return "xywh (COCO style)" # c > a and d > b — ambiguous between xyxy and xywh. # Use image dimensions to disambiguate when available. if image_size is not None: img_w, img_h = image_size # If interpreting as xywh, right edge = a + c; if that overshoots the # image while c alone fits, the format is more likely xyxy. xywh_exceeds = (a + c > img_w * 1.05) or (b + d > img_h * 1.05) xyxy_exceeds = (c > img_w * 1.05) or (d > img_h * 1.05) if xywh_exceeds and not xyxy_exceeds: return "xyxy (Pascal VOC style)" if xyxy_exceeds and not xywh_exceeds: return "xywh (COCO style)" if is_normalized: return "xyxy_normalized" return "xyxy (Pascal VOC style)" def _extract_image_size(row: Dict) -> Tuple[int, int] | None: """Try to extract (width, height) from the image column returned by Datasets Server.""" for col in ("image", "img", "picture", "photo"): img = row.get(col) if isinstance(img, dict): w = img.get("width") h = img.get("height") if isinstance(w, (int, float)) and isinstance(h, (int, float)): return (int(w), int(h)) return None def analyze_annotations(sample_rows: List[Dict], annotation_cols: List[str]) -> Dict[str, Any]: """Analyze annotation structure from sample rows""" if not annotation_cols: return {"found": False} annotation_col = annotation_cols[0] annotations_info = { "found": True, "column": annotation_col, "sample_structures": [], "bbox_formats": [], "categories_found": [], "avg_objects_per_image": 0, "max_objects": 0, "min_objects": float('inf'), } total_objects = 0 valid_samples = 0 for row in sample_rows: ann = row["row"].get(annotation_col) if not ann: continue valid_samples += 1 image_size = _extract_image_size(row["row"]) # Check if it's a list of annotations or a dict if isinstance(ann, dict): # COCO-style or structured annotation sample_structure = { "type": "dict", "keys": list(ann.keys()) } # Check for bounding boxes if "bbox" in ann or "bboxes" in ann: bbox_key = "bbox" if "bbox" in ann else "bboxes" bboxes = ann[bbox_key] if isinstance(bboxes, list) and len(bboxes) > 0: if isinstance(bboxes[0], list): # Multiple bboxes num_objects = len(bboxes) total_objects += num_objects annotations_info["max_objects"] = max(annotations_info["max_objects"], num_objects) annotations_info["min_objects"] = min(annotations_info["min_objects"], num_objects) # Analyze first bbox format bbox_format = detect_bbox_format(bboxes[0], image_size) annotations_info["bbox_formats"].append(bbox_format) else: # Single bbox total_objects += 1 annotations_info["max_objects"] = max(annotations_info["max_objects"], 1) annotations_info["min_objects"] = min(annotations_info["min_objects"], 1) bbox_format = detect_bbox_format(bboxes, image_size) annotations_info["bbox_formats"].append(bbox_format) # Check for categories/classes for key in ["category", "categories", "label", "labels", "class", "classes", "category_id"]: if key in ann: cats = ann[key] if isinstance(cats, list): annotations_info["categories_found"].extend([str(c) for c in cats]) else: annotations_info["categories_found"].append(str(cats)) annotations_info["sample_structures"].append(sample_structure) elif isinstance(ann, list): # List of annotation dicts sample_structure = { "type": "list", "length": len(ann), "item_type": type(ann[0]).__name__ if ann else None } if ann and isinstance(ann[0], dict): sample_structure["item_keys"] = list(ann[0].keys()) # Count objects num_objects = len(ann) total_objects += num_objects annotations_info["max_objects"] = max(annotations_info["max_objects"], num_objects) annotations_info["min_objects"] = min(annotations_info["min_objects"], num_objects) # Check first annotation first_ann = ann[0] if "bbox" in first_ann: bbox_format = detect_bbox_format(first_ann["bbox"], image_size) annotations_info["bbox_formats"].append(bbox_format) # Check for categories for key in ["category", "label", "class", "category_id"]: if key in first_ann: for item in ann: if key in item: annotations_info["categories_found"].append(str(item[key])) annotations_info["sample_structures"].append(sample_structure) if valid_samples > 0: annotations_info["avg_objects_per_image"] = round(total_objects / valid_samples, 2) if annotations_info["min_objects"] == float('inf'): annotations_info["min_objects"] = 0 # Get unique categories annotations_info["categories_found"] = list(set(annotations_info["categories_found"])) annotations_info["num_classes"] = len(annotations_info["categories_found"]) # Get most common bbox format if annotations_info["bbox_formats"]: from collections import Counter format_counts = Counter(annotations_info["bbox_formats"]) annotations_info["primary_bbox_format"] = format_counts.most_common(1)[0][0] return annotations_info def check_image_classification_compatibility(columns: List[str], sample_rows: List[Dict], features: List[Dict]) -> Dict[str, Any]: """Check image classification dataset compatibility""" image_cols = find_columns(columns, ["image", "img", "picture", "photo"]) has_image = len(image_cols) > 0 label_cols = find_columns(columns, ["label", "labels", "class", "fine_label", "coarse_label"]) has_label = len(label_cols) > 0 label_info: Dict[str, Any] = {"found": has_label} if has_label: label_col = label_cols[0] label_info["column"] = label_col # Detect whether label is ClassLabel (int with names) or plain int/string for f in features: if f.get("name") == label_col: ftype = f.get("type", "") if isinstance(ftype, dict) and ftype.get("_type") == "ClassLabel": label_info["type"] = "ClassLabel" names = ftype.get("names", []) label_info["num_classes"] = len(names) label_info["class_names"] = names[:20] if len(names) > 20: label_info["class_names_truncated"] = True elif isinstance(ftype, dict) and ftype.get("dtype") in ("int64", "int32", "int8"): label_info["type"] = "int" elif isinstance(ftype, dict) and ftype.get("dtype") == "string": label_info["type"] = "string" break # Discover unique labels from samples if ClassLabel info wasn't in features if "num_classes" not in label_info: unique = set() for row in sample_rows: val = row["row"].get(label_col) if val is not None: unique.add(val) label_info["sample_unique_labels"] = sorted(unique, key=str)[:20] label_info["sample_unique_count"] = len(unique) ready = has_image and has_label return { "ready": ready, "has_image": has_image, "image_columns": image_cols, "has_label": has_label, "label_columns": label_cols, "label_info": label_info, } def check_object_detection_compatibility(columns: List[str], sample_rows: List[Dict]) -> Dict[str, Any]: """Check object detection dataset compatibility""" # Find image column image_cols = find_columns(columns, ["image", "img", "picture", "photo"]) has_image = len(image_cols) > 0 # Find annotation columns annotation_cols = find_columns(columns, ["objects", "annotations", "ann", "bbox", "bboxes", "detection"]) has_annotations = len(annotation_cols) > 0 # Analyze annotations annotations_info = analyze_annotations(sample_rows, annotation_cols) if has_annotations else {"found": False} # Check for separate bbox and category columns bbox_cols = find_columns(columns, ["bbox", "bboxes", "boxes"]) category_cols = find_columns(columns, ["category", "label", "class", "categories", "labels", "classes"]) # Determine readiness ready = has_image and (has_annotations or (len(bbox_cols) > 0 and len(category_cols) > 0)) return { "ready": ready, "has_image": has_image, "image_columns": image_cols, "has_annotations": has_annotations, "annotation_columns": annotation_cols, "separate_bbox_columns": bbox_cols, "separate_category_columns": category_cols, "annotations_info": annotations_info, } def check_sam_segmentation_compatibility(columns: List[str], sample_rows: List[Dict], features: List[Dict]) -> Dict[str, Any]: """Check SAM/SAM2 segmentation dataset compatibility. A valid SAM segmentation dataset needs: - An image column - A mask column (binary ground-truth segmentation mask) - A prompt: either a bbox prompt or point prompt (in a JSON prompt column, or dedicated columns) """ image_cols = find_columns(columns, ["image", "img", "picture", "photo"]) has_image = len(image_cols) > 0 mask_cols = find_columns(columns, ["mask", "segmentation", "alpha", "matte"]) has_mask = len(mask_cols) > 0 prompt_cols = find_columns(columns, ["prompt"]) bbox_cols = [c for c in columns if c in ("bbox", "bboxes", "box", "boxes")] point_cols = [c for c in columns if c in ("point", "points", "input_point", "input_points")] prompt_info: Dict[str, Any] = { "has_prompt": False, "prompt_type": None, "source": None, "bbox_valid": None, } # Try JSON prompt column first if prompt_cols: for row in sample_rows: raw = row["row"].get(prompt_cols[0]) if raw is None: continue parsed = raw if isinstance(raw, dict) else _try_json(raw) if parsed is None: continue if isinstance(parsed, dict): if "bbox" in parsed or "box" in parsed: prompt_info["has_prompt"] = True prompt_info["prompt_type"] = "bbox" prompt_info["source"] = f"JSON column '{prompt_cols[0]}'" bbox = parsed.get("bbox") or parsed.get("box") prompt_info["bbox_valid"] = _validate_bbox(bbox, _extract_image_size(row["row"])) break elif "point" in parsed or "points" in parsed: prompt_info["has_prompt"] = True prompt_info["prompt_type"] = "point" prompt_info["source"] = f"JSON column '{prompt_cols[0]}'" break if not prompt_info["has_prompt"] and bbox_cols: prompt_info["has_prompt"] = True prompt_info["prompt_type"] = "bbox" prompt_info["source"] = f"column '{bbox_cols[0]}'" for row in sample_rows: bbox = row["row"].get(bbox_cols[0]) if bbox is not None: prompt_info["bbox_valid"] = _validate_bbox(bbox, _extract_image_size(row["row"])) break if not prompt_info["has_prompt"] and point_cols: prompt_info["has_prompt"] = True prompt_info["prompt_type"] = "point" prompt_info["source"] = f"column '{point_cols[0]}'" ready = has_image and has_mask and prompt_info["has_prompt"] return { "ready": ready, "has_image": has_image, "image_columns": image_cols, "has_mask": has_mask, "mask_columns": mask_cols, "prompt_columns": prompt_cols, "bbox_columns": bbox_cols, "point_columns": point_cols, "prompt_info": prompt_info, } def _try_json(value) -> Any: if not isinstance(value, str): return None try: return json.loads(value) except (json.JSONDecodeError, TypeError): return None def _validate_bbox(bbox, image_size=None) -> Dict[str, Any]: """Validate a single bounding box and return diagnostics.""" result: Dict[str, Any] = {"valid": False} if not isinstance(bbox, (list, tuple)): result["error"] = "bbox is not a list" return result if len(bbox) != 4: result["error"] = f"expected 4 values, got {len(bbox)}" return result try: vals = [float(v) for v in bbox] except (TypeError, ValueError): result["error"] = "non-numeric values" return result if not all(math.isfinite(v) for v in vals): result["error"] = "contains non-finite values" return result x0, y0, x1, y1 = vals if x1 <= x0 or y1 <= y0: if vals[2] > 0 and vals[3] > 0: result["format_hint"] = "likely xywh" else: result["error"] = "degenerate bbox (zero or negative area)" return result else: result["format_hint"] = "likely xyxy" if image_size is not None: img_w, img_h = image_size if any(v > max(img_w, img_h) * 1.5 for v in vals): result["warning"] = "coordinates exceed image bounds" result["valid"] = True result["values"] = vals return result def generate_mapping_code(info: Dict[str, Any]) -> str: """Generate mapping code if needed""" if info["ready"]: ann_info = info["annotations_info"] if not ann_info.get("found"): return None # Check if format conversion is needed ann_col = ann_info.get("column") bbox_format = ann_info.get("primary_bbox_format", "unknown") if "coco" in bbox_format.lower() or "xywh" in bbox_format.lower(): # Already COCO format return f"""# Dataset appears to be in COCO format (xywh) # Image column: {info['image_columns'][0] if info['image_columns'] else 'image'} # Annotation column: {ann_col} # Use directly with transformers object detection models""" elif "xyxy" in bbox_format.lower(): # Need to convert from XYXY to XYWH return f"""# Convert from XYXY (Pascal VOC) to XYWH (COCO) format def convert_to_coco_format(example): annotations = example['{ann_col}'] if isinstance(annotations, list): for ann in annotations: if 'bbox' in ann: x_min, y_min, x_max, y_max = ann['bbox'] ann['bbox'] = [x_min, y_min, x_max - x_min, y_max - y_min] elif isinstance(annotations, dict) and 'bbox' in annotations: bbox = annotations['bbox'] if isinstance(bbox, list) and len(bbox) > 0 and isinstance(bbox[0], list): for i, box in enumerate(bbox): x_min, y_min, x_max, y_max = box bbox[i] = [x_min, y_min, x_max - x_min, y_max - y_min] return example dataset = dataset.map(convert_to_coco_format)""" elif not info["ready"]: # Need to create annotations structure if info["separate_bbox_columns"] and info["separate_category_columns"]: bbox_col = info["separate_bbox_columns"][0] cat_col = info["separate_category_columns"][0] return f"""# Combine separate bbox and category columns def create_annotations(example): bboxes = example['{bbox_col}'] categories = example['{cat_col}'] if not isinstance(bboxes, list): bboxes = [bboxes] if not isinstance(categories, list): categories = [categories] annotations = [] for bbox, cat in zip(bboxes, categories): annotations.append({{'bbox': bbox, 'category': cat}}) example['objects'] = annotations return example dataset = dataset.map(cre ... ``` ### `scripts/estimate_cost.py` ``` #!/usr/bin/env python3 # /// script # requires-python = ">=3.10" # dependencies = [] # /// """ Estimate training time and cost for vision model training jobs on Hugging Face Jobs. Usage: uv run estimate_cost.py --model ustc-community/dfine-small-coco --dataset cppe-5 --hardware t4-small uv run estimate_cost.py --model PekingU/rtdetr_v2_r50vd --dataset-size 5000 --hardware t4-small --epochs 30 uv run estimate_cost.py --model google/vit-base-patch16-224-in21k --dataset ethz/food101 --hardware t4-small --epochs 3 """ import argparse HARDWARE_COSTS = { "t4-small": 0.40, "t4-medium": 0.60, "l4x1": 0.80, "l4x4": 3.80, "a10g-small": 1.00, "a10g-large": 1.50, "a10g-largex2": 3.00, "a10g-largex4": 5.00, "l40sx1": 1.80, "l40sx4": 8.30, "a100-large": 2.50, "a100x4": 10.00, } # Vision model sizes in millions of parameters MODEL_PARAMS_M = { # Object detection "dfine-small": 10.4, "dfine-large": 31.4, "dfine-xlarge": 63.5, "rtdetr_v2_r18vd": 20.2, "rtdetr_v2_r50vd": 43.0, "rtdetr_v2_r101vd": 76.0, "detr-resnet-50": 41.3, "detr-resnet-101": 60.2, "yolos-small": 30.7, "yolos-tiny": 6.5, # Image classification "mobilenetv3_small": 2.5, "mobilevit_s": 5.6, "resnet50": 25.6, "vit_base_patch16": 86.6, # SAM / SAM2 segmentation "sam-vit-base": 93.7, "sam-vit-large": 312.3, "sam-vit-huge": 641.1, "sam2.1-hiera-tiny": 38.9, "sam2.1-hiera-small": 46.0, "sam2.1-hiera-base-plus": 80.8, "sam2.1-hiera-large": 224.4, } KNOWN_DATASETS = { # Object detection "cppe-5": 1000, "merve/license-plate": 6180, # Image classification "ethz/food101": 75750, # SAM segmentation "merve/MicroMat-mini": 240, } def extract_model_params(model_name: str) -> float: """Extract model size in millions of parameters from the model name.""" name_lower = model_name.lower() for key, params in MODEL_PARAMS_M.items(): if key.lower() in name_lower: return params return 30.0 # reasonable default for vision models def estimate_training_time(model_params_m: float, dataset_size: int, epochs: int, image_size: int, batch_size: int, hardware: str) -> float: """Estimate training time in hours for vision model training.""" # Steps per epoch steps_per_epoch = dataset_size / batch_size # empirical calibration values base_secs_per_step = 0.8 model_factor = (model_params_m / 30.0) ** 0.6 image_factor = (image_size / 640.0) ** 2 batch_factor = (batch_size / 8.0) ** 0.7 secs_per_step = base_secs_per_step * model_factor * image_factor * batch_factor hardware_multipliers = { "t4-small": 2.0, "t4-medium": 2.0, "l4x1": 1.2, "l4x4": 0.5, "a10g-small": 1.0, "a10g-large": 1.0, "a10g-largex2": 0.6, "a10g-largex4": 0.4, "l40sx1": 0.7, "l40sx4": 0.25, "a100-large": 0.5, "a100x4": 0.2, } multiplier = hardware_multipliers.get(hardware, 1.0) total_steps = steps_per_epoch * epochs total_secs = total_steps * secs_per_step * multiplier # Add overhead: model loading (~2 min), eval per epoch (~10% of training), Hub push (~3 min) eval_overhead = total_secs * 0.10 fixed_overhead = 5 * 60 # 5 minutes total_secs += eval_overhead + fixed_overhead return total_secs / 3600 def parse_args(): parser = argparse.ArgumentParser(description="Estimate training cost for vision model training jobs") parser.add_argument("--model", required=True, help="Model name (e.g., 'ustc-community/dfine-small-coco' or 'detr-resnet-50')") parser.add_argument("--dataset", default=None, help="Dataset name (for known size lookup)") parser.add_argument("--hardware", required=True, choices=HARDWARE_COSTS.keys(), help="Hardware flavor") parser.add_argument("--dataset-size", type=int, default=None, help="Number of training images (overrides dataset lookup)") parser.add_argument("--epochs", type=int, default=30, help="Number of training epochs (default: 30)") parser.add_argument("--image-size", type=int, default=640, help="Image square size in pixels (default: 640)") parser.add_argument("--batch-size", type=int, default=8, help="Per-device batch size (default: 8)") return parser.parse_args() def main(): args = parse_args() model_params = extract_model_params(args.model) print(f"Model: {args.model} (~{model_params:.1f}M parameters)") if args.dataset_size: dataset_size = args.dataset_size elif args.dataset and args.dataset in KNOWN_DATASETS: dataset_size = KNOWN_DATASETS[args.dataset] elif args.dataset: print(f"Unknown dataset '{args.dataset}', defaulting to 1000 images.") print(f"Use --dataset-size to specify the exact count.") dataset_size = 1000 else: dataset_size = 1000 print(f"Dataset: {args.dataset or 'custom'} (~{dataset_size} images)") print(f"Epochs: {args.epochs}") print(f"Image size: {args.image_size}px") print(f"Batch size: {args.batch_size}") print(f"Hardware: {args.hardware} (${HARDWARE_COSTS[args.hardware]:.2f}/hr)") print() estimated_hours = estimate_training_time( model_params, dataset_size, args.epochs, args.image_size, args.batch_size, args.hardware ) estimated_cost = estimated_hours * HARDWARE_COSTS[args.hardware] recommended_timeout = estimated_hours * 1.3 # 30% buffer print(f"Estimated training time: {estimated_hours:.1f} hours") print(f"Estimated cost: ${estimated_cost:.2f}") print(f"Recommended timeout: {recommended_timeout:.1f}h (with 30% buffer)") print() if estimated_hours > 6: print("Warning: Long training time. Consider:") print(" - Reducing epochs or image size") print(" - Using --max_train_samples for a test run first") print(" - Upgrading hardware") print() if model_params > 50 and args.hardware in ("t4-small", "t4-medium"): print("Warning: Large model on T4. If you hit OOM:") print(" - Reduce batch size (try 4, then 2)") print(" - Reduce image size (try 480)") print(" - Upgrade to l4x1 or a10g-small") print() timeout_str = f"{recommended_timeout:.0f}h" timeout_secs = int(recommended_timeout * 3600) print(f"Example job configuration (MCP tool):") print(f""" hf_jobs("uv", {{ "script": "scripts/object_detection_training.py", "script_args": [ "--model_name_or_path", "{args.model}", "--dataset_name", "{args.dataset or 'your-dataset'}", "--image_square_size", "{args.image_size}", "--num_train_epochs", "{args.epochs}", "--per_device_train_batch_size", "{args.batch_size}", "--push_to_hub", "--do_train", "--do_eval" ], "flavor": "{args.hardware}", "timeout": "{timeout_str}", "secrets": {{"HF_TOKEN": "$HF_TOKEN"}} }}) """) print(f"Example job configuration (Python API):") print(f""" api.run_uv_job( script="scripts/object_detection_training.py", script_args=[...], flavor="{args.hardware}", timeout={timeout_secs}, secrets={{"HF_TOKEN": get_token()}}, ) """) if __name__ == "__main__": main() ``` ### `scripts/image_classification_training.py` ``` # /// script # dependencies = [ # "transformers>=5.2.0", # "accelerate>=1.1.0", # "timm", # "datasets>=4.0", # "evaluate", # "scikit-learn", # "torchvision", # "trackio", # "huggingface_hub", # ] # /// """Fine-tuning any Transformers or timm model supported by AutoModelForImageClassification using the Trainer API.""" import logging import os import sys from dataclasses import dataclass, field from functools import partial from typing import Any import evaluate import numpy as np import torch from datasets import load_dataset from torchvision.transforms import ( CenterCrop, Compose, Normalize, RandomHorizontalFlip, RandomResizedCrop, Resize, ToTensor, ) import trackio import transformers from transformers import ( AutoConfig, AutoImageProcessor, AutoModelForImageClassification, DefaultDataCollator, HfArgumentParser, Trainer, TrainingArguments, ) from transformers.trainer import EvalPrediction from transformers.utils import check_min_version from transformers.utils.versions import require_version logger = logging.getLogger(__name__) check_min_version("4.57.0.dev0") require_version("datasets>=2.0.0") @dataclass class DataTrainingArguments: dataset_name: str = field( default="ethz/food101", metadata={"help": "Name of a dataset from the Hub."}, ) dataset_config_name: str | None = field( default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}, ) train_val_split: float | None = field( default=0.15, metadata={"help": "Fraction to split off of train for validation (used only when no validation split exists)."}, ) max_train_samples: int | None = field( default=None, metadata={"help": "Truncate training set to this many samples (for debugging / quick tests)."}, ) max_eval_samples: int | None = field( default=None, metadata={"help": "Truncate evaluation set to this many samples."}, ) image_column_name: str = field( default="image", metadata={"help": "The column name for images in the dataset."}, ) label_column_name: str = field( default="label", metadata={"help": "The column name for labels in the dataset."}, ) @dataclass class ModelArguments: model_name_or_path: str = field( default="timm/mobilenetv3_small_100.lamb_in1k", metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models."}, ) config_name: str | None = field( default=None, metadata={"help": "Pretrained config name or path if not the same as model_name."}, ) cache_dir: str | None = field( default=None, metadata={"help": "Where to store pretrained models downloaded from the Hub."}, ) model_revision: str = field( default="main", metadata={"help": "The specific model version to use (branch, tag, or commit id)."}, ) image_processor_name: str | None = field( default=None, metadata={"help": "Name or path of image processor config."}, ) ignore_mismatched_sizes: bool = field( default=True, metadata={"help": "Allow loading weights when num_labels differs from pretrained checkpoint."}, ) token: str | None = field( default=None, metadata={"help": "Auth token for private models / datasets."}, ) trust_remote_code: bool = field( default=False, metadata={"help": "Whether to trust remote code from Hub repos."}, ) def build_transforms(image_processor, is_training: bool): """Build torchvision transforms from the image processor's config.""" if hasattr(image_processor, "size"): size = image_processor.size if "shortest_edge" in size: img_size = size["shortest_edge"] elif "height" in size and "width" in size: img_size = (size["height"], size["width"]) else: img_size = 224 else: img_size = 224 if hasattr(image_processor, "image_mean") and image_processor.image_mean: normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std) else: normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if is_training: return Compose([ RandomResizedCrop(img_size), RandomHorizontalFlip(), ToTensor(), normalize, ]) else: if isinstance(img_size, int): resize_size = int(img_size / 0.875) # standard 87.5% center crop ratio else: resize_size = tuple(int(s / 0.875) for s in img_size) return Compose([ Resize(resize_size), CenterCrop(img_size), ToTensor(), normalize, ]) def main(): parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1])) else: model_args, data_args, training_args = parser.parse_args_into_dataclasses() # --- Hub authentication --- from huggingface_hub import login hf_token = os.environ.get("HF_TOKEN") or os.environ.get("hfjob") if hf_token: login(token=hf_token) training_args.hub_token = hf_token logger.info("Logged in to Hugging Face Hub") elif training_args.push_to_hub: logger.warning("HF_TOKEN not found in environment. Hub push will likely fail.") # --- Trackio --- trackio.init(project=training_args.output_dir, name=training_args.run_name) # --- Logging --- logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", handlers=[logging.StreamHandler(sys.stdout)], ) if training_args.should_log: transformers.utils.logging.set_verbosity_info() log_level = training_args.get_process_log_level() logger.setLevel(log_level) transformers.utils.logging.set_verbosity(log_level) transformers.utils.logging.enable_default_handler() transformers.utils.logging.enable_explicit_format() logger.warning( f"Process rank: {training_args.local_process_index}, device: {training_args.device}, " f"n_gpu: {training_args.n_gpu}, distributed training: " f"{training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" ) logger.info(f"Training/evaluation parameters {training_args}") # --- Load dataset --- dataset = load_dataset( data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, trust_remote_code=model_args.trust_remote_code, ) # --- Resolve label column --- label_col = data_args.label_column_name if label_col not in dataset["train"].column_names: candidates = [c for c in dataset["train"].column_names if c in ("label", "labels", "class", "fine_label")] if candidates: label_col = candidates[0] logger.info(f"Label column '{data_args.label_column_name}' not found, using '{label_col}'") else: raise ValueError( f"Label column '{data_args.label_column_name}' not found. " f"Available columns: {dataset['train'].column_names}" ) # --- Discover labels --- label_feature = dataset["train"].features[label_col] if hasattr(label_feature, "names"): label_names = label_feature.names else: unique_labels = sorted(set(dataset["train"][label_col])) if all(isinstance(l, str) for l in unique_labels): label_names = unique_labels else: label_names = [str(l) for l in unique_labels] num_labels = len(label_names) id2label = dict(enumerate(label_names)) label2id = {v: k for k, v in id2label.items()} logger.info(f"Number of classes: {num_labels}") # --- Remap string labels to int if needed --- sample_label = dataset["train"][0][label_col] if isinstance(sample_label, str): logger.info("Remapping string labels to integer IDs") for split_name in list(dataset.keys()): dataset[split_name] = dataset[split_name].map( lambda ex: {label_col: label2id[ex[label_col]]}, ) # --- Shuffle + Train/val split --- dataset["train"] = dataset["train"].shuffle(seed=training_args.seed) data_args.train_val_split = None if "validation" in dataset else data_args.train_val_split if isinstance(data_args.train_val_split, float) and data_args.train_val_split > 0.0: split = dataset["train"].train_test_split(data_args.train_val_split, seed=training_args.seed) dataset["train"] = split["train"] dataset["validation"] = split["test"] # --- Truncate --- if data_args.max_train_samples is not None: max_train = min(data_args.max_train_samples, len(dataset["train"])) dataset["train"] = dataset["train"].select(range(max_train)) logger.info(f"Truncated training set to {max_train} samples") if data_args.max_eval_samples is not None and "validation" in dataset: max_eval = min(data_args.max_eval_samples, len(dataset["validation"])) dataset["validation"] = dataset["validation"].select(range(max_eval)) logger.info(f"Truncated validation set to {max_eval} samples") # --- Load model & image processor --- common_pretrained_args = { "cache_dir": model_args.cache_dir, "revision": model_args.model_revision, "token": model_args.token, "trust_remote_code": model_args.trust_remote_code, } config = AutoConfig.from_pretrained( model_args.config_name or model_args.model_name_or_path, num_labels=num_labels, label2id=label2id, id2label=id2label, **common_pretrained_args, ) model = AutoModelForImageClassification.from_pretrained( model_args.model_name_or_path, config=config, ignore_mismatched_sizes=model_args.ignore_mismatched_sizes, **common_pretrained_args, ) image_processor = AutoImageProcessor.from_pretrained( model_args.image_processor_name or model_args.model_name_or_path, **common_pretrained_args, ) # --- Build transforms --- train_transforms = build_transforms(image_processor, is_training=True) val_transforms = build_transforms(image_processor, is_training=False) image_col = data_args.image_column_name def preprocess_train(examples): return { "pixel_values": [train_transforms(img.convert("RGB")) for img in examples[image_col]], "labels": examples[label_col], } def preprocess_val(examples): return { "pixel_values": [val_transforms(img.convert("RGB")) for img in examples[image_col]], "labels": examples[label_col], } dataset["train"].set_transform(preprocess_train) if "validation" in dataset: dataset["validation"].set_transform(preprocess_val) if "test" in dataset: dataset["test"].set_transform(preprocess_val) # --- Metrics --- accuracy_metric = evaluate.load("accuracy") def compute_metrics(eval_pred: EvalPrediction): predictions = np.argmax(eval_pred.predictions, axis=1) return accuracy_metric.compute(predictions=predictions, references=eval_pred.label_ids) # --- Trainer --- eval_dataset = None if training_args.do_eval: if "validation" in dataset: eval_dataset = dataset["validation"] elif "test" in dataset: eval_dataset = dataset["test"] trainer = Trainer( model=model, args=training_args, train_dataset=dataset["train"] if training_args.do_train else None, eval_dataset=eval_dataset, processing_class=image_processor, data_collator=DefaultDataCollator(), compute_metrics=compute_metrics, ) # --- Train --- if training_args.do_train: train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint) trainer.save_model() trainer.log_metrics("train", train_result.metrics) trainer.save_metrics("train", train_result.metrics) trainer.save_state() # --- Evaluate --- if training_args.do_eval: test_dataset = dataset.get("test", dataset.get("validation")) test_prefix = "test" if "test" in dataset else "eval" if test_dataset is not None: metrics = trainer.evaluate(eval_dataset=test_dataset, metric_key_prefix=test_prefix) trainer.log_metrics(test_prefix, metrics) trainer.save_metrics(test_prefix, metrics) trackio.finish() # --- Push to Hub --- kwargs = { "finetuned_from": model_args.model_name_or_path, "dataset": data_args.dataset_name, "tags": ["image-classification", "vision"], } if training_args.push_to_hub: trainer.push_to_hub(**kwargs) else: trainer.create_model_card(**kwargs) if __name__ == "__main__": main() ``` ### `scripts/object_detection_training.py` ``` # /// script # dependencies = [ # "transformers>=5.2.0", # "accelerate>=1.1.0", # "albumentations >= 1.4.16", # "timm", # "datasets>=4.0", # "torchmetrics", # "pycocotools", # "trackio", # "huggingface_hub", # ] # /// """Finetuning any 🤗 Transformers model supported by AutoModelForObjectDetection for object detection leveraging the Trainer API.""" import logging import math import os import sys from collections.abc import Mapping from dataclasses import dataclass, field from functools import partial from typing import Any import albumentations as A import numpy as np import torch from datasets import load_dataset from torchmetrics.detection.mean_ap import MeanAveragePrecision import trackio import transformers from transformers import ( AutoConfig, AutoImageProcessor, AutoModelForObjectDetection, HfArgumentParser, Trainer, TrainingArguments, ) from transformers.image_processing_utils import BatchFeature from transformers.image_transforms import center_to_corners_format from transformers.trainer import EvalPrediction from transformers.utils import check_min_version from transformers.utils.versions import require_version logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers is not installed. Remove at your own risks. check_min_version("4.57.0.dev0") require_version("datasets>=2.0.0", "To fix: pip install -r examples/pytorch/object-detection/requirements.txt") @dataclass class ModelOutput: logits: torch.Tensor pred_boxes: torch.Tensor def format_image_annotations_as_coco( image_id: str, categories: list[int], areas: list[float], bboxes: list[tuple[float]] ) -> dict: """Format one set of image annotations to the COCO format Args: image_id (str): image id. e.g. "0001" categories (list[int]): list of categories/class labels corresponding to provided bounding boxes areas (list[float]): list of corresponding areas to provided bounding boxes bboxes (list[tuple[float]]): list of bounding boxes provided in COCO format ([center_x, center_y, width, height] in absolute coordinates) Returns: dict: { "image_id": image id, "annotations": list of formatted annotations } """ annotations = [] for category, area, bbox in zip(categories, areas, bboxes): formatted_annotation = { "image_id": image_id, "category_id": category, "iscrowd": 0, "area": area, "bbox": list(bbox), } annotations.append(formatted_annotation) return { "image_id": image_id, "annotations": annotations, } def detect_bbox_format_from_samples(dataset, image_col="image", objects_col="objects", num_samples=50): """ Detect whether bboxes are xyxy (Pascal VOC) or xywh (COCO) by checking bbox coordinates against image dimensions. The correct format interpretation should keep bboxes within image bounds. """ exceeds_if_xywh = 0 exceeds_if_xyxy = 0 total = 0 for example in dataset.select(range(min(num_samples, len(dataset)))): img_w, img_h = example[image_col].size for bbox in example[objects_col]["bbox"]: if len(bbox) != 4: continue a, b, c, d = float(bbox[0]), float(bbox[1]), float(bbox[2]), float(bbox[3]) total += 1 # If 3rd < 1st or 4th < 2nd, can't be xyxy (x_max must exceed x_min) if c < a or d < b: return "xywh" # xywh: right/bottom edge = origin + size; exceeding image → wrong format if a + c > img_w * 1.05: exceeds_if_xywh += 1 if b + d > img_h * 1.05: exceeds_if_xywh += 1 # xyxy: right/bottom edge = coordinate itself if c > img_w * 1.05: exceeds_if_xyxy += 1 if d > img_h * 1.05: exceeds_if_xyxy += 1 if total == 0: return "xywh" fmt = "xyxy" if exceeds_if_xywh > exceeds_if_xyxy else "xywh" logger.info( f"Detected bbox format: {fmt} (checked {total} bboxes from {min(num_samples, len(dataset))} images)" ) return fmt def sanitize_dataset(dataset, bbox_format="xywh", image_col="image", objects_col="objects"): """ Validate bboxes, convert xyxy→xywh if needed, clip to image bounds, and remove entries with non-finite values, non-positive dimensions, or degenerate area (<1 px). Drops images with no remaining valid bboxes. """ convert_xyxy = bbox_format == "xyxy" def _validate(example): img_w, img_h = example[image_col].size objects = example[objects_col] bboxes = objects["bbox"] n = len(bboxes) valid_indices = [] converted_bboxes = [] for i, bbox in enumerate(bboxes): if len(bbox) != 4: continue vals = [float(v) for v in bbox] if not all(math.isfinite(v) for v in vals): continue if convert_xyxy: x_min, y_min, x_max, y_max = vals w, h = x_max - x_min, y_max - y_min else: x_min, y_min, w, h = vals if w <= 0 or h <= 0: continue x_min, y_min = max(0.0, x_min), max(0.0, y_min) if x_min >= img_w or y_min >= img_h: continue w = min(w, img_w - x_min) h = min(h, img_h - y_min) if w * h < 1.0: continue valid_indices.append(i) converted_bboxes.append([x_min, y_min, w, h]) # Rebuild objects dict, filtering all list-valued fields by valid_indices new_objects = {} for key, value in objects.items(): if key == "bbox": new_objects["bbox"] = converted_bboxes elif isinstance(value, list) and len(value) == n: new_objects[key] = [value[j] for j in valid_indices] else: new_objects[key] = value if "area" not in new_objects or len(new_objects.get("area", [])) != len(converted_bboxes): new_objects["area"] = [b[2] * b[3] for b in converted_bboxes] example[objects_col] = new_objects return example before = len(dataset) dataset = dataset.map(_validate) dataset = dataset.filter(lambda ex: len(ex[objects_col]["bbox"]) > 0) after = len(dataset) if before != after: logger.warning(f"Dropped {before - after}/{before} images with no valid bboxes after sanitization") logger.info(f"Bbox sanitization complete: {after} images with valid bboxes remain") return dataset def convert_bbox_yolo_to_pascal(boxes: torch.Tensor, image_size: tuple[int, int]) -> torch.Tensor: """ Convert bounding boxes from YOLO format (x_center, y_center, width, height) in range [0, 1] to Pascal VOC format (x_min, y_min, x_max, y_max) in absolute coordinates. Args: boxes (torch.Tensor): Bounding boxes in YOLO format image_size (tuple[int, int]): Image size in format (height, width) Returns: torch.Tensor: Bounding boxes in Pascal VOC format (x_min, y_min, x_max, y_max) """ # convert center to corners format boxes = center_to_corners_format(boxes) if isinstance(image_size, torch.Tensor): image_size = image_size.tolist() elif isinstance(image_size, np.ndarray): image_size = image_size.tolist() height, width = image_size boxes = boxes * torch.tensor([[width, height, width, height]]) return boxes def augment_and_transform_batch( examples: Mapping[str, Any], transform: A.Compose, image_processor: AutoImageProcessor, return_pixel_mask: bool = False, ) -> BatchFeature: """Apply augmentations and format annotations in COCO format for object detection task""" images = [] annotations = [] image_ids = examples["image_id"] if "image_id" in examples else range(len(examples["image"])) for image_id, image, objects in zip(image_ids, examples["image"], examples["objects"]): image = np.array(image.convert("RGB")) # Filter invalid bboxes before augmentation (safety net after sanitize_dataset) bboxes = objects["bbox"] categories = objects["category"] areas = objects["area"] valid = [ (b, c, a) for b, c, a in zip(bboxes, categories, areas) if len(b) == 4 and b[2] > 0 and b[3] > 0 and b[0] >= 0 and b[1] >= 0 ] if valid: bboxes, categories, areas = zip(*valid) else: bboxes, categories, areas = [], [], [] # apply augmentations output = transform(image=image, bboxes=list(bboxes), category=list(categories)) images.append(output["image"]) # format annotations in COCO format (recompute areas from post-augmentation bboxes) post_areas = [b[2] * b[3] for b in output["bboxes"]] if output["bboxes"] else [] formatted_annotations = format_image_annotations_as_coco( image_id, output["category"], post_areas, output["bboxes"] ) annotations.append(formatted_annotations) # Apply the image processor transformations: resizing, rescaling, normalization result = image_processor(images=images, annotations=annotations, return_tensors="pt") if not return_pixel_mask: result.pop("pixel_mask", None) return result def collate_fn(batch: list[BatchFeature]) -> Mapping[str, torch.Tensor | list[Any]]: data = {} data["pixel_values"] = torch.stack([x["pixel_values"] for x in batch]) data["labels"] = [x["labels"] for x in batch] if "pixel_mask" in batch[0]: data["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch]) return data @torch.no_grad() def compute_metrics( evaluation_results: EvalPrediction, image_processor: AutoImageProcessor, threshold: float = 0.0, id2label: Mapping[int, str] | None = None, ) -> Mapping[str, float]: """ Compute mean average mAP, mAR and their variants for the object detection task. Args: evaluation_results (EvalPrediction): Predictions and targets from evaluation. threshold (float, optional): Threshold to filter predicted boxes by confidence. Defaults to 0.0. id2label (Optional[dict], optional): Mapping from class id to class name. Defaults to None. Returns: Mapping[str, float]: Metrics in a form of dictionary {: } """ predictions, targets = evaluation_results.predictions, evaluation_results.label_ids # For metric computation we need to provide: # - targets in a form of list of dictionaries with keys "boxes", "labels" # - predictions in a form of list of dictionaries with keys "boxes", "scores", "labels" image_sizes = [] post_processed_targets = [] post_processed_predictions = [] # Collect targets in the required format for metric computation for batch in targets: # collect image sizes, we will need them for predictions post processing batch_image_sizes = torch.tensor([x["orig_size"] for x in batch]) image_sizes.append(batch_image_sizes) # collect targets in the required format for metric computation # boxes were converted to YOLO format needed for model training # here we will convert them to Pascal VOC format (x_min, y_min, x_max, y_max) for image_target in batch: boxes = torch.tensor(image_target["boxes"]) boxes = convert_bbox_yolo_to_pascal(boxes, image_target["orig_size"]) labels = torch.tensor(image_target["class_labels"]) post_processed_targets.append({"boxes": boxes, "labels": labels}) # Collect predictions in the required format for metric computation, # model produce boxes in YOLO format, then image_processor convert them to Pascal VOC format for batch, target_sizes in zip(predictions, image_sizes): batch_logits, batch_boxes = batch[1], batch[2] output = ModelOutput(logits=torch.tensor(batch_logits), pred_boxes=torch.tensor(batch_boxes)) post_processed_output = image_processor.post_process_object_detection( output, threshold=threshold, target_sizes=target_sizes ) post_processed_predictions.extend(post_processed_output) # Compute metrics metric = MeanAveragePrecision(box_format="xyxy", class_metrics=True) metric.update(post_processed_predictions, post_processed_targets) metrics = metric.compute() # Replace list of per class metrics with separate metric for each class classes = metrics.pop("classes") map_per_class = metrics.pop("map_per_class") mar_100_per_class = metrics.pop("mar_100_per_class") # Single-class datasets return 0-d scalar tensors; make them iterable if classes.dim() == 0: classes = classes.unsqueeze(0) map_per_class = map_per_class.unsqueeze(0) mar_100_per_class = mar_100_per_class.unsqueeze(0) for class_id, class_map, class_mar in zip(classes, map_per_class, mar_100_per_class): class_name = id2label[class_id.item()] if id2label is not None else class_id.item() metrics[f"map_{class_name}"] = class_map metrics[f"mar_100_{class_name}"] = class_mar metrics = {k: round(v.item(), 4) for k, v in metrics.items()} return metrics @dataclass class DataTrainingArguments: """ Arguments pertaining to what data we are going to input our model for training and eval. Using `HfArgumentParser` we can turn this class into argparse arguments to be able to specify them on the command line. """ dataset_name: str = field( default="cppe-5", metadata={ "help": "Name of a dataset from the hub (could be your own, possibly private dataset hosted on the hub)." }, ) dataset_config_name: str | None = field( default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} ) train_val_split: float | None = field( default=0.15, metadata={"help": "Percent to split off of train for validation."} ) image_square_size: int | None = field( default=600, metadata={"help": "Image longest size will be resized to this value, then image will be padded to square."}, ) max_train_samples: int | None = field( default=None, metadata={ "help": ( "For debugging purposes or quicker training, truncate the number of training examples to this " "value if set." ) }, ) max_eval_samples: int | None = field( default=None, metadata={ "help": ( "For debugging purposes or quicker training, truncate the number of evaluation examples to this " "value if set." ) }, ) use_fast: bool | None = field( default=True, metadata={"help": "Use a fast torchvision-base image processor if it is supported for a given model."}, ) @dataclass class ModelArguments: """ Arguments pertaining to which model/config/tokenizer we are going to fine-tune from. """ model_name_or_path: str = field( default="facebook/detr-resnet-50", metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}, ) config_name: str | None = field( default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"} ) cache_dir: str | None = field( default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"} ) model_revision: str = field( default="main", metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}, ) image_processor_name: str = field(default=None, metadata={"help": "Name or path of preprocessor config."}) ignore_mismatched_sizes: bool = field( default=True, metadata={ "help": "Whether or not to raise an error if some of the weights from the checkpoint do not have the same size as the weights of the model (if for instance, you are instantiating a model with 10 labels from a checkpoint with 3 labels)." }, ) token: str = field( default=None, metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) trust_remote_code: bool = field( default=False, metadata={ "help": ( "Whether to trust the execution of code from datasets/models defined on the Hub." " This option should only be set to `True` for repositories you trust and in which you have read the" " code, as it will execute code present on the Hub on your local machine." ) }, ) def main(): parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1])) else: model_args, data_args, training_args = parser.parse_args_into_dataclasses() from huggingface_hub import login hf_token = os.environ.get("HF_TOKEN") or os.environ.get("hfjob") if hf_token: login(token=hf_token) training_args.hub_token = hf_token logger.info("Logged in to Hugging Face Hub") elif training_args.push_to_hub: logger.warning("HF_TOKEN not found in environment. Hub push will likely fail.") # Initialize Trackio for real-time experiment tracking trackio.init(project=training_args.output_dir, name=training_args.run_name) logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", handlers=[logging.StreamHandler(sys.stdout)], ) if training_args.should_log: # The default of training_args.log_level is passive, so we set log level at info here to have that default. transformers.utils.logging.set_verbosity_info() log_level = training_args.get_process_log_level() logger.setLevel(log_level) transformers.utils.logging.set_verbosity(log_level) transformers.utils.logging.enable_default_handler() transformers.utils.logging.enable_explicit_format() # Log on each process the small summary: logger.warning( f"Process rank: {training_args.local_process_index}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" ) logger.info(f"Training/evaluation parameters {training_args}") dataset = load_dataset( data_args.dataset_name, cache_dir=model_args.cache_dir, trust_remote_code=model_args.trust_remote_code ) bbox_format = detect_bbox_format_from_samples(dataset["train"]) if bbox_format == "xyxy": logger.info("Converting bboxes from xyxy (Pascal VOC) → xywh (COCO) format across all splits") for split_name in list(dataset.keys()): dataset[split_name] = sanitize_dataset(dataset[split_name], bbox_format=bbox_format) for split_name in list(dataset.keys()): if "image_id" not in dataset[split_name].column_names: dataset[split_name] = dataset[split_name].add_column( "image_id", list(range(len(dataset[split_name]))) ) dataset[ ... ``` ### `scripts/sam_segmentation_training.py` ``` # /// script # dependencies = [ # "transformers>=5.2.0", # "accelerate>=1.1.0", # "datasets>=4.0", # "torchvision", # "monai", # "trackio", # "huggingface_hub", # ] # /// """Fine-tune SAM or SAM2 for segmentation using bounding-box or point prompts with the HF Trainer API.""" import json import logging import math import os import sys from dataclasses import dataclass, field from typing import Any import numpy as np import torch import torch.nn.functional as F from datasets import load_dataset from torch.utils.data import Dataset import monai import trackio import transformers from transformers import ( HfArgumentParser, Trainer, TrainingArguments, ) from transformers.utils import check_min_version logger = logging.getLogger(__name__) check_min_version("4.57.0.dev0") # --------------------------------------------------------------------------- # Dataset wrapper # --------------------------------------------------------------------------- class SAMSegmentationDataset(Dataset): """Wraps a HF dataset into the format expected by SAM/SAM2 processors. Each sample must contain an image, a binary mask, and a prompt (bbox or point). Prompts are read from a JSON-encoded ``prompt`` column or from dedicated ``bbox`` / ``point`` columns. """ def __init__(self, dataset, processor, prompt_type: str, image_col: str, mask_col: str, prompt_col: str | None, bbox_col: str | None, point_col: str | None): self.dataset = dataset self.processor = processor self.prompt_type = prompt_type self.image_col = image_col self.mask_col = mask_col self.prompt_col = prompt_col self.bbox_col = bbox_col self.point_col = point_col def __len__(self): return len(self.dataset) def _extract_prompt(self, item): if self.prompt_col and self.prompt_col in item: raw = item[self.prompt_col] parsed = json.loads(raw) if isinstance(raw, str) else raw if self.prompt_type == "bbox": return parsed.get("bbox") or parsed.get("box") return parsed.get("point") or parsed.get("points") if self.prompt_type == "bbox" and self.bbox_col: return item[self.bbox_col] if self.prompt_type == "point" and self.point_col: return item[self.point_col] raise ValueError("Could not extract prompt from sample") def __getitem__(self, idx): item = self.dataset[idx] image = item[self.image_col] prompt = self._extract_prompt(item) if self.prompt_type == "bbox": inputs = self.processor(image, input_boxes=[[prompt]], return_tensors="pt") else: if isinstance(prompt[0], (int, float)): prompt = [prompt] inputs = self.processor(image, input_points=[[prompt]], return_tensors="pt") mask = np.array(item[self.mask_col]) if mask.ndim == 3: mask = mask[:, :, 0] inputs["labels"] = (mask > 0).astype(np.float32) inputs["original_image_size"] = torch.tensor(image.size[::-1]) return inputs def collate_fn(batch): pixel_values = torch.cat([item["pixel_values"] for item in batch], dim=0) original_sizes = torch.stack([item["original_sizes"] for item in batch]) original_image_size = torch.stack([item["original_image_size"] for item in batch]) has_boxes = "input_boxes" in batch[0] has_points = "input_points" in batch[0] labels = torch.cat( [ F.interpolate( torch.as_tensor(x["labels"]).unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode="nearest", ) for x in batch ], dim=0, ).long() result = { "pixel_values": pixel_values, "original_sizes": original_sizes, "labels": labels, "original_image_size": original_image_size, "multimask_output": False, } if has_boxes: result["input_boxes"] = torch.cat([item["input_boxes"] for item in batch], dim=0) if has_points: result["input_points"] = torch.cat([item["input_points"] for item in batch], dim=0) if "input_labels" in batch[0]: result["input_labels"] = torch.cat([item["input_labels"] for item in batch], dim=0) return result # --------------------------------------------------------------------------- # Custom loss (SAM/SAM2 don't compute loss in forward()) # --------------------------------------------------------------------------- seg_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction="mean") def compute_loss(outputs, labels, num_items_in_batch=None): predicted_masks = outputs.pred_masks.squeeze(1) return seg_loss(predicted_masks, labels.float()) # --------------------------------------------------------------------------- # CLI arguments # --------------------------------------------------------------------------- @dataclass class DataTrainingArguments: dataset_name: str = field( default="merve/MicroMat-mini", metadata={"help": "Hub dataset ID."}, ) dataset_config_name: str | None = field( default=None, metadata={"help": "Dataset config name."}, ) train_val_split: float | None = field( default=0.1, metadata={"help": "Fraction to split off for validation (used when no validation split exists)."}, ) max_train_samples: int | None = field( default=None, metadata={"help": "Truncate training set (for quick tests)."}, ) max_eval_samples: int | None = field( default=None, metadata={"help": "Truncate evaluation set."}, ) image_column_name: str = field( default="image", metadata={"help": "Column containing PIL images."}, ) mask_column_name: str = field( default="mask", metadata={"help": "Column containing ground-truth binary masks."}, ) prompt_column_name: str | None = field( default="prompt", metadata={"help": "Column with JSON-encoded prompt (bbox/point). Set to '' to disable."}, ) bbox_column_name: str | None = field( default=None, metadata={"help": "Column with bbox prompt ([x0,y0,x1,y1]). Used when prompt_column_name is unset."}, ) point_column_name: str | None = field( default=None, metadata={"help": "Column with point prompt ([x,y] or [[x,y],...]). Used when prompt_column_name is unset."}, ) prompt_type: str = field( default="bbox", metadata={"help": "Prompt type: 'bbox' or 'point'."}, ) @dataclass class ModelArguments: model_name_or_path: str = field( default="facebook/sam2.1-hiera-small", metadata={"help": "Pretrained SAM/SAM2 model identifier."}, ) cache_dir: str | None = field(default=None, metadata={"help": "Cache directory."}) model_revision: str = field(default="main", metadata={"help": "Model revision."}) token: str | None = field(default=None, metadata={"help": "Auth token."}) trust_remote_code: bool = field(default=False, metadata={"help": "Trust remote code."}) freeze_vision_encoder: bool = field( default=True, metadata={"help": "Freeze vision encoder weights."}, ) freeze_prompt_encoder: bool = field( default=True, metadata={"help": "Freeze prompt encoder weights."}, ) # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) parser.set_defaults(per_device_train_batch_size=4, num_train_epochs=30) if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): model_args, data_args, training_args = parser.parse_json_file( json_file=os.path.abspath(sys.argv[1]) ) else: model_args, data_args, training_args = parser.parse_args_into_dataclasses() from huggingface_hub import login hf_token = os.environ.get("HF_TOKEN") or os.environ.get("hfjob") if hf_token: login(token=hf_token) training_args.hub_token = hf_token logger.info("Logged in to Hugging Face Hub") elif training_args.push_to_hub: logger.warning("HF_TOKEN not found in environment. Hub push will likely fail.") trackio.init(project=training_args.output_dir, name=training_args.run_name) logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", handlers=[logging.StreamHandler(sys.stdout)], ) if training_args.should_log: transformers.utils.logging.set_verbosity_info() log_level = training_args.get_process_log_level() logger.setLevel(log_level) transformers.utils.logging.set_verbosity(log_level) transformers.utils.logging.enable_default_handler() transformers.utils.logging.enable_explicit_format() logger.info(f"Training/evaluation parameters {training_args}") # ---- Load dataset ---- dataset = load_dataset( data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, trust_remote_code=model_args.trust_remote_code, ) if "train" not in dataset: if len(dataset.keys()) == 1: only_split = list(dataset.keys())[0] dataset[only_split] = dataset[only_split].shuffle(seed=training_args.seed) dataset = dataset[only_split].train_test_split(test_size=data_args.train_val_split or 0.1) dataset = {"train": dataset["train"], "validation": dataset["test"]} else: raise ValueError(f"No 'train' split found. Available: {list(dataset.keys())}") elif "validation" not in dataset and "test" not in dataset: dataset["train"] = dataset["train"].shuffle(seed=training_args.seed) split = dataset["train"].train_test_split( test_size=data_args.train_val_split or 0.1, seed=training_args.seed ) dataset["train"] = split["train"] dataset["validation"] = split["test"] if data_args.max_train_samples is not None: n = min(data_args.max_train_samples, len(dataset["train"])) dataset["train"] = dataset["train"].select(range(n)) logger.info(f"Truncated training set to {n} samples") eval_key = "validation" if "validation" in dataset else "test" if data_args.max_eval_samples is not None and eval_key in dataset: n = min(data_args.max_eval_samples, len(dataset[eval_key])) dataset[eval_key] = dataset[eval_key].select(range(n)) logger.info(f"Truncated eval set to {n} samples") # ---- Detect model family (SAM vs SAM2) and load processor/model ---- model_id = model_args.model_name_or_path.lower() is_sam2 = "sam2" in model_id if is_sam2: from transformers import Sam2Processor, Sam2Model processor = Sam2Processor.from_pretrained(model_args.model_name_or_path) model = Sam2Model.from_pretrained(model_args.model_name_or_path) else: from transformers import SamProcessor, SamModel processor = SamProcessor.from_pretrained(model_args.model_name_or_path) model = SamModel.from_pretrained(model_args.model_name_or_path) if model_args.freeze_vision_encoder: for name, param in model.named_parameters(): if name.startswith("vision_encoder"): param.requires_grad_(False) if model_args.freeze_prompt_encoder: for name, param in model.named_parameters(): if name.startswith("prompt_encoder"): param.requires_grad_(False) trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) total = sum(p.numel() for p in model.parameters()) logger.info(f"Trainable params: {trainable:,} / {total:,} ({100 * trainable / total:.1f}%)") # ---- Build datasets ---- prompt_col = data_args.prompt_column_name if data_args.prompt_column_name else None ds_kwargs = dict( processor=processor, prompt_type=data_args.prompt_type, image_col=data_args.image_column_name, mask_col=data_args.mask_column_name, prompt_col=prompt_col, bbox_col=data_args.bbox_column_name, point_col=data_args.point_column_name, ) train_dataset = SAMSegmentationDataset(dataset=dataset["train"], **ds_kwargs) eval_dataset = None if eval_key in dataset: eval_dataset = SAMSegmentationDataset(dataset=dataset[eval_key], **ds_kwargs) # ---- Train ---- trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, data_collator=collate_fn, compute_loss_func=compute_loss, ) if training_args.do_train: train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint) trainer.save_model() trainer.log_metrics("train", train_result.metrics) trainer.save_metrics("train", train_result.metrics) trainer.save_state() if training_args.do_eval and eval_dataset is not None: metrics = trainer.evaluate() trainer.log_metrics("eval", metrics) trainer.save_metrics("eval", metrics) trackio.finish() kwargs = { "finetuned_from": model_args.model_name_or_path, "dataset": data_args.dataset_name, "tags": ["image-segmentation", "vision", "sam"], } if training_args.push_to_hub: trainer.push_to_hub(**kwargs) else: trainer.create_model_card(**kwargs) if __name__ == "__main__": main() ```