Upload train.py

train.py

@@ -0,0 +1,132 @@
+"""
+Fine-tune a lightweight image classifier on rice grain (seed) images.
+Dataset: nateraw/rice-image-dataset (5 rice varieties, 75K images)
+Model: microsoft/resnet-18 (lightweight, ~11M params)
+
+Usage (with Hugging Face Jobs):
+  HF_MODEL_REPO=chaosbee997/rice-seed-classifier python train.py
+
+Or submit via hf_jobs with a10g-large / t4-small GPU.
+"""
+import os
+import numpy as np
+from datasets import load_dataset
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+    TrainingArguments,
+    Trainer,
+)
+from transformers import DefaultDataCollator
+from PIL import Image
+import evaluate
+
+# ============= CONFIG =============
+DATASET_NAME = "nateraw/rice-image-dataset"
+MODEL_NAME = "microsoft/resnet-18"
+OUTPUT_REPO = os.environ.get("HF_MODEL_REPO", "chaosbee997/rice-seed-classifier")
+
+# ============= LOAD DATASET =============
+print(f"Loading dataset: {DATASET_NAME}")
+ds = load_dataset(DATASET_NAME)
+print(ds)
+
+split_ds = ds["train"].train_test_split(test_size=0.15, stratify_by_column="label", seed=42)
+train_ds = split_ds["train"]
+val_ds = split_ds["test"]
+
+labels = ds["train"].features["label"].names
+num_labels = len(labels)
+label2id = {label: i for i, label in enumerate(labels)}
+id2label = {i: label for i, label in enumerate(labels)}
+
+print(f"Classes ({num_labels}): {labels}")
+print(f"Train: {len(train_ds)} | Val: {len(val_ds)}")
+
+# ============= LOAD PROCESSOR & MODEL =============
+print(f"Loading model: {MODEL_NAME}")
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+model = AutoModelForImageClassification.from_pretrained(
+    MODEL_NAME,
+    num_labels=num_labels,
+    id2label=id2label,
+    label2id=label2id,
+    ignore_mismatched_sizes=True,
+)
+
+# ============= PREPROCESS =============
+def transform(example_batch):
+    images = []
+    for img in example_batch["image"]:
+        if isinstance(img, Image.Image):
+            if img.mode != "RGB":
+                img = img.convert("RGB")
+            images.append(img)
+        else:
+            img = Image.fromarray(np.array(img)).convert("RGB")
+            images.append(img)
+    inputs = processor(images, return_tensors="pt")
+    inputs["labels"] = example_batch["label"]
+    return inputs
+
+train_ds.set_transform(transform)
+val_ds.set_transform(transform)
+
+# ============= METRICS =============
+accuracy = evaluate.load("accuracy")
+f1 = evaluate.load("f1")
+
+def compute_metrics(eval_pred):
+    predictions, labels = eval_pred
+    preds = np.argmax(predictions, axis=1)
+    acc = accuracy.compute(predictions=preds, references=labels)
+    f1_score = f1.compute(predictions=preds, references=labels, average="weighted")
+    return {"accuracy": acc["accuracy"], "f1": f1_score["f1"]}
+
+# ============= TRAINING ARGS =============
+args = TrainingArguments(
+    output_dir="/tmp/rice-seed-classifier",
+    remove_unused_columns=False,
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    learning_rate=5e-5,
+    per_device_train_batch_size=64,
+    per_device_eval_batch_size=64,
+    num_train_epochs=5,
+    warmup_ratio=0.1,
+    logging_strategy="steps",
+    logging_steps=50,
+    logging_first_step=True,
+    disable_tqdm=True,
+    load_best_model_at_end=True,
+    metric_for_best_model="accuracy",
+    seed=42,
+    push_to_hub=True,
+    hub_model_id=OUTPUT_REPO,
+    report_to="trackio",
+    run_name="rice_resnet18_lr5e-5_bs64",
+    project="grain-classification",
+    trackio_space_id=os.environ.get("TRACKIO_SPACE_ID", "chaosbee997/mlintern-grain"),
+)
+
+# ============= TRAINER =============
+trainer = Trainer(
+    model=model,
+    args=args,
+    train_dataset=train_ds,
+    eval_dataset=val_ds,
+    compute_metrics=compute_metrics,
+    data_collator=DefaultDataCollator(),
+    tokenizer=processor,
+)
+
+print("Starting training...")
+trainer.train()
+
+print("Evaluating...")
+metrics = trainer.evaluate()
+print(metrics)
+
+print("Pushing to hub...")
+trainer.push_to_hub()
+print("Done!")