tribunal-kit 9.2.3 → 9.2.4

package/skills/huggingface-vision-trainer/scripts/estimate_cost.py

@@ -0,0 +1,217 @@
+#!/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()

package/skills/huggingface-vision-trainer/scripts/image_classification_training.py

@@ -0,0 +1,383 @@
+# /// 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()