# 🚀 Google Colab 训练指南

本指南教你如何在 **免费的 Google Colab T4 GPU** 上训练籽粒分类模型。

---

## 第一步：打开 Colab

1. 访问 [colab.research.google.com](https://colab.research.google.com)
2. 点击 **文件 → 新建笔记本**

---

## 第二步：切换到 GPU 环境

点击菜单栏：
```
Runtime → Change runtime type → Hardware accelerator → T4 GPU
```
然后点击 **Save**，并 **重新连接**（点击左上角重新连接按钮）。

---

## 第三步：逐 Cell 运行代码

将下面的代码按顺序复制到 Colab 的每个 Cell 中运行：

### Cell 1：检查 GPU
```python
!nvidia-smi
```
> 应该能看到 `Tesla T4` 信息。如果看不到，回到第二步确认 GPU 已启用。

---

### Cell 2：安装依赖
```python
!pip install -q transformers datasets accelerate evaluate pillow huggingface_hub
print("✅ 依赖安装完成")
```

---

### Cell 3：导入库
```python
import os
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer,
    DefaultDataCollator,
)
from PIL import Image
import evaluate
from huggingface_hub import notebook_login
```

---

### Cell 4：登录 Hugging Face（必须！）
```python
notebook_login()
```
> 运行后会弹出一个输入框，要求输入 **Access Token**。
>
> 🔑 **如何获取 Token？**
> 1. 打开 [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)
> 2. 点击 `New token`
> 3. 选择 **Write** 权限
> 4. 复制 token 粘贴到 Colab 弹窗中

---

### Cell 5：加载数据集
```python
DATASET_NAME = "nateraw/rice-image-dataset"
print(f"Loading dataset: {DATASET_NAME}...")
ds = load_dataset(DATASET_NAME)
print(ds)

labels = ds["train"].features["label"].names
num_labels = len(labels)
print(f"\n📋 类别数量: {num_labels}")
print(f"📋 类别列表: {labels}")
```

---

### Cell 6：划分训练/验证集
```python
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"]

print(f"训练集: {len(train_ds)} 张")
print(f"验证集: {len(val_ds)} 张")
```

---

### Cell 7：加载模型
```python
MODEL_NAME = "microsoft/resnet-18"

processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

label2id = {label: i for i, label in enumerate(labels)}
id2label = {i: label for i, label in enumerate(labels)}

model = AutoModelForImageClassification.from_pretrained(
    MODEL_NAME,
    num_labels=num_labels,
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes=True,
)

print(f"✅ 模型加载完成")
print(f"参数总量: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")
```

---

### Cell 8：数据预处理
```python
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)
print("✅ 数据预处理设置完成")
```

---

### Cell 9：评估指标
```python
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"]}

print("✅ 评估指标定义完成")
```

---

### Cell 10：训练参数配置
```python
OUTPUT_REPO = "chaosbee997/rice-seed-classifier"  # ← 改成你的仓库名

args = TrainingArguments(
    output_dir="/content/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,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    seed=42,
    push_to_hub=True,
    hub_model_id=OUTPUT_REPO,
    report_to="none",
)

print(f"训练配置:")
print(f"  输出仓库: {OUTPUT_REPO}")
print(f"  学习率: {args.learning_rate}")
print(f"  Batch size: {args.per_device_train_batch_size}")
print(f"  Epochs: {args.num_train_epochs}")
```

---

### Cell 11：开始训练 🚀
```python
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    compute_metrics=compute_metrics,
    data_collator=DefaultDataCollator(),
    tokenizer=processor,
)

print("开始训练...\n" + "="*50)
trainer.train()
```
> ⏱️ 训练时间：T4 GPU 上约 **15-30 分钟**（5 epochs，63K训练样本）
>
> 📊 预期结果：准确率 > **95%**

---

### Cell 12：评估并上传模型
```python
metrics = trainer.evaluate()
print("\n最终评估结果:")
print("="*50)
for k, v in metrics.items():
    if isinstance(v, float):
        print(f"  {k}: {v:.4f}")
    else:
        print(f"  {k}: {v}")

print("\n正在上传到 Hugging Face Hub...")
trainer.push_to_hub()
print(f"\n✅ 上传成功!")
print(f"🔗 模型地址: https://huggingface.co/{OUTPUT_REPO}")
```

---

### Cell 13：推理测试
```python
from transformers import pipeline
import matplotlib.pyplot as plt

classifier = pipeline("image-classification", model=OUTPUT_REPO)

# 从数据集取一张图测试
test_image = ds["train"][100]["image"]

plt.figure(figsize=(5, 5))
plt.imshow(test_image)
plt.axis("off")
plt.title("测试图像")
plt.show()

results = classifier(test_image)
print("\n🔮 预测结果:")
for r in results:
    print(f"  {r['label']}: {r['score']*100:.2f}%")
```

---

## ⚡ 训练完成后

你的模型会自动上传到：
- **https://huggingface.co/chaosbee997/rice-seed-classifier**

任何人都可以用一行代码加载你的模型：
```python
from transformers import pipeline
classifier = pipeline("image-classification", model="chaosbee997/rice-seed-classifier")
```

---

## 🌽 扩展到花生/玉米/小麦等其他作物

如果你想训练包含 **花生、玉米、小麦、水稻** 的通用籽粒分类模型：

1. **收集图像**：每种作物建立一个文件夹，例如：
   ```
   crop_seeds/
   ├── peanut/
   │   ├── img1.jpg
   │   └── img2.jpg
   ├── corn/
   │   ├── img1.jpg
   │   └── img2.jpg
   ├── wheat/
   │   ├── img1.jpg
   │   └── img2.jpg
   └── rice/
       ├── img1.jpg
       └── img2.jpg
   ```

2. **创建数据集**：
   ```python
   from datasets import load_dataset
   ds = load_dataset("imagefolder", data_dir="/path/to/crop_seeds")
   ds.push_to_hub("yourname/crop-seeds")
   ```

3. **修改 Cell 5**：把 `DATASET_NAME` 改成你的数据集名称

4. **重新运行 Cell 5-12**，其余代码完全不用改！

---

## 🆘 常见问题

**Q: 训练时报 `CUDA out of memory`？**  
A: 减小 batch size：把 `per_device_train_batch_size=64` 改成 `32` 或 `16`。

**Q: 没有 Hugging Face 账号？**  
A: 免费注册：[huggingface.co/join](https://huggingface.co/join)

**Q: 训练时间太长？**  
A: 可以减少 epoch 数（`num_train_epochs=3`），或换更小的模型（`microsoft/resnet-34` 更大，`google/mobilenet_v2_1.0_224` 更轻）。