{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "cells": [ { "cell_type": "markdown", "id": "title_md", "metadata": {}, "source": [ "# RF-DETR: Training on a Custom Dataset\n", "\n", "**Series:** RF-DETR Detailed Tutorial | **Notebook 16d of 5**\n", "\n", "---\n", "\n", "## What we'll cover\n", "\n", "1. Dataset format requirements (COCO JSON)\n", "2. Build a synthetic robot manipulation dataset from scratch\n", "3. Implement a custom `Dataset` class for RF-DETR\n", "4. Configure the model for a new number of classes\n", "5. Training loop with Hungarian matching loss\n", "6. Evaluation: mAP computation\n", "7. Fine-tuning strategies: frozen backbone vs full fine-tune\n", "8. Save & reload fine-tuned checkpoint\n", "\n", "---\n", "\n", "## Dataset Format\n", "\n", "RF-DETR (via HuggingFace Transformers) expects **COCO JSON format**:\n", "\n", "```json\n", "{\n", " \"images\": [{\"id\": 1, \"file_name\": \"img001.jpg\", \"height\": 480, \"width\": 640}],\n", " \"annotations\": [\n", " {\"id\": 1, \"image_id\": 1, \"category_id\": 1,\n", " \"bbox\": [x, y, width, height], // COCO format: top-left + size\n", " \"area\": 1200, \"iscrowd\": 0}\n", " ],\n", " \"categories\": [{\"id\": 1, \"name\": \"robot_arm\"}, {\"id\": 2, \"name\": \"object\"}]\n", "}\n", "```" ] }, { "cell_type": "markdown", "id": "dataset_md", "metadata": {}, "source": [ "## 1. Create a Synthetic Robot Manipulation Dataset\n", "\n", "We'll generate a small synthetic dataset with:\n", "- **3 classes**: `robot_arm`, `gripper`, `target_object`\n", "- **50 training images**, **10 validation images**\n", "- Synthetic images with colored rectangles (simulating robot parts)\n", "\n", "This lets us demonstrate the full training pipeline without needing real labeled data." ] }, { "cell_type": "code", "execution_count": null, "id": "create_dataset", "metadata": {}, "outputs": [], "source": [ "import os\n", "import json\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import matplotlib.patches as patches\n", "from PIL import Image, ImageDraw\n", "import random\n", "import torch\n", "import warnings\n", "warnings.filterwarnings('ignore')\n", "\n", "random.seed(42)\n", "np.random.seed(42)\n", "\n", "# ── Dataset configuration ─────────────────────────────────────────────────────\n", "DATASET_DIR = '/tmp/robot_detection_dataset'\n", "CLASSES = ['robot_arm', 'gripper', 'target_object']\n", "IMG_W, IMG_H = 640, 480\n", "N_TRAIN, N_VAL = 50, 10\n", "\n", "# Class colors for synthetic images\n", "CLASS_COLORS = {\n", " 'robot_arm': (70, 130, 180), # steel blue\n", " 'gripper': (220, 100, 50), # orange-red\n", " 'target_object': (50, 180, 80), # green\n", "}\n", "\n", "os.makedirs(f'{DATASET_DIR}/train', exist_ok=True)\n", "os.makedirs(f'{DATASET_DIR}/val', exist_ok=True)\n", "\n", "def generate_synthetic_image(img_id, split):\n", " \"\"\"Generate a synthetic robot scene with bounding box annotations.\"\"\"\n", " # Background: gradient gray\n", " bg = np.ones((IMG_H, IMG_W, 3), dtype=np.uint8) * 40\n", " bg += np.random.randint(0, 20, bg.shape, dtype=np.uint8)\n", " img = Image.fromarray(bg)\n", " draw = ImageDraw.Draw(img)\n", " \n", " annotations = []\n", " ann_id_base = img_id * 10\n", " \n", " # Always draw robot_arm (large rectangle)\n", " arm_x = random.randint(50, 300)\n", " arm_y = random.randint(50, 200)\n", " arm_w = random.randint(80, 200)\n", " arm_h = random.randint(30, 80)\n", " draw.rectangle([arm_x, arm_y, arm_x+arm_w, arm_y+arm_h],\n", " fill=CLASS_COLORS['robot_arm'], outline=(255,255,255), width=2)\n", " annotations.append({\n", " 'id': ann_id_base + 1,\n", " 'image_id': img_id,\n", " 'category_id': 1, # robot_arm\n", " 'bbox': [arm_x, arm_y, arm_w, arm_h],\n", " 'area': arm_w * arm_h,\n", " 'iscrowd': 0\n", " })\n", " \n", " # Draw gripper (small rectangle at end of arm)\n", " grip_x = arm_x + arm_w - 10\n", " grip_y = arm_y + arm_h // 4\n", " grip_w = random.randint(20, 50)\n", " grip_h = random.randint(20, 50)\n", " draw.rectangle([grip_x, grip_y, grip_x+grip_w, grip_y+grip_h],\n", " fill=CLASS_COLORS['gripper'], outline=(255,255,255), width=2)\n", " annotations.append({\n", " 'id': ann_id_base + 2,\n", " 'image_id': img_id,\n", " 'category_id': 2, # gripper\n", " 'bbox': [grip_x, grip_y, grip_w, grip_h],\n", " 'area': grip_w * grip_h,\n", " 'iscrowd': 0\n", " })\n", " \n", " # Randomly add target object\n", " if random.random() > 0.3:\n", " obj_x = random.randint(300, 550)\n", " obj_y = random.randint(200, 380)\n", " obj_w = random.randint(30, 80)\n", " obj_h = random.randint(30, 80)\n", " draw.ellipse([obj_x, obj_y, obj_x+obj_w, obj_y+obj_h],\n", " fill=CLASS_COLORS['target_object'], outline=(255,255,255), width=2)\n", " annotations.append({\n", " 'id': ann_id_base + 3,\n", " 'image_id': img_id,\n", " 'category_id': 3, # target_object\n", " 'bbox': [obj_x, obj_y, obj_w, obj_h],\n", " 'area': obj_w * obj_h,\n", " 'iscrowd': 0\n", " })\n", " \n", " # Add noise\n", " noise = np.random.randint(0, 15, (IMG_H, IMG_W, 3), dtype=np.uint8)\n", " img = Image.fromarray(np.clip(np.array(img).astype(int) + noise, 0, 255).astype(np.uint8))\n", " \n", " fname = f'img_{img_id:04d}.jpg'\n", " img.save(f'{DATASET_DIR}/{split}/{fname}', quality=90)\n", " \n", " return {\n", " 'id': img_id,\n", " 'file_name': fname,\n", " 'height': IMG_H,\n", " 'width': IMG_W\n", " }, annotations\n", "\n", "# Generate train set\n", "train_images, train_annotations = [], []\n", "for i in range(1, N_TRAIN + 1):\n", " img_info, anns = generate_synthetic_image(i, 'train')\n", " train_images.append(img_info)\n", " train_annotations.extend(anns)\n", "\n", "# Generate val set\n", "val_images, val_annotations = [], []\n", "for i in range(N_TRAIN + 1, N_TRAIN + N_VAL + 1):\n", " img_info, anns = generate_synthetic_image(i, 'val')\n", " val_images.append(img_info)\n", " val_annotations.extend(anns)\n", "\n", "# COCO categories\n", "categories = [{'id': i+1, 'name': cls, 'supercategory': 'robot'}\n", " for i, cls in enumerate(CLASSES)]\n", "\n", "# Save COCO JSON\n", "for split, images, annotations in [('train', train_images, train_annotations),\n", " ('val', val_images, val_annotations)]:\n", " coco_data = {'images': images, 'annotations': annotations, 'categories': categories}\n", " with open(f'{DATASET_DIR}/{split}_annotations.json', 'w') as f:\n", " json.dump(coco_data, f, indent=2)\n", "\n", "print(f'✅ Dataset created at {DATASET_DIR}')\n", "print(f' Train: {len(train_images)} images, {len(train_annotations)} annotations')\n", "print(f' Val: {len(val_images)} images, {len(val_annotations)} annotations')\n", "print(f' Classes: {CLASSES}')" ] }, { "cell_type": "code", "execution_count": null, "id": "visualize_dataset", "metadata": {}, "outputs": [], "source": [ "# ── Visualize sample training images ─────────────────────────────────────────\n", "with open(f'{DATASET_DIR}/train_annotations.json') as f:\n", " train_coco = json.load(f)\n", "\n", "# Build lookup\n", "img_id_to_anns = {}\n", "for ann in train_coco['annotations']:\n", " img_id_to_anns.setdefault(ann['image_id'], []).append(ann)\n", "cat_id_to_name = {cat['id']: cat['name'] for cat in train_coco['categories']}\n", "cat_colors = {'robot_arm': '#4682b4', 'gripper': '#dc6432', 'target_object': '#32b450'}\n", "\n", "fig, axes = plt.subplots(2, 4, figsize=(20, 10))\n", "fig.patch.set_facecolor('#0d1117')\n", "fig.suptitle('Synthetic Robot Manipulation Dataset — Sample Images',\n", " color='white', fontsize=13, fontweight='bold')\n", "\n", "for ax, img_info in zip(axes.flat, train_coco['images'][:8]):\n", " ax.set_facecolor('#0d1117')\n", " img = Image.open(f'{DATASET_DIR}/train/{img_info[\"file_name\"]}')\n", " ax.imshow(np.array(img))\n", " ax.set_title(f'Image {img_info[\"id\"]}', color='white', fontsize=9)\n", " ax.axis('off')\n", " \n", " for ann in img_id_to_anns.get(img_info['id'], []):\n", " x, y, w, h = ann['bbox']\n", " cls_name = cat_id_to_name[ann['category_id']]\n", " color = cat_colors[cls_name]\n", " rect = patches.Rectangle((x, y), w, h,\n", " linewidth=2, edgecolor=color, facecolor='none')\n", " ax.add_patch(rect)\n", " ax.text(x, y - 3, cls_name, color='white', fontsize=7,\n", " bbox=dict(facecolor=color, alpha=0.8, pad=1, edgecolor='none'))\n", "\n", "plt.tight_layout()\n", "plt.savefig('rfdetr_dataset_samples.png', dpi=100, bbox_inches='tight', facecolor='#0d1117')\n", "plt.show()\n", "print('Dataset samples saved.')" ] }, { "cell_type": "markdown", "id": "dataset_class_md", "metadata": {}, "source": [ "## 2. Custom Dataset Class\n", "\n", "We implement a PyTorch `Dataset` that:\n", "1. Loads images from disk\n", "2. Applies the RF-DETR image processor\n", "3. Returns targets in the format expected by the model\n", "\n", "**Target format for RF-DETR:**\n", "```python\n", "{\n", " 'class_labels': torch.LongTensor([cls_id, ...]),\n", " 'boxes': torch.FloatTensor([[cx, cy, w, h], ...]) # normalized [0,1]\n", "}\n", "```" ] }, { "cell_type": "code", "execution_count": null, "id": "dataset_class", "metadata": {}, "outputs": [], "source": [ "import torch\n", "from torch.utils.data import Dataset, DataLoader\n", "from transformers import AutoImageProcessor\n", "\n", "class RobotDetectionDataset(Dataset):\n", " \"\"\"\n", " Custom dataset for RF-DETR training.\n", " Loads COCO-format annotations and applies the RF-DETR image processor.\n", " \"\"\"\n", " \n", " def __init__(self, image_dir, annotation_file, processor, augment=False):\n", " self.image_dir = image_dir\n", " self.processor = processor\n", " self.augment = augment\n", " \n", " with open(annotation_file) as f:\n", " coco = json.load(f)\n", " \n", " self.images = coco['images']\n", " self.categories = {cat['id']: cat['name'] for cat in coco['categories']}\n", " \n", " # Build image_id → annotations lookup\n", " self.img_id_to_anns = {}\n", " for ann in coco['annotations']:\n", " self.img_id_to_anns.setdefault(ann['image_id'], []).append(ann)\n", " \n", " def __len__(self):\n", " return len(self.images)\n", " \n", " def __getitem__(self, idx):\n", " img_info = self.images[idx]\n", " img_id = img_info['id']\n", " img_w = img_info['width']\n", " img_h = img_info['height']\n", " \n", " # Load image\n", " img_path = os.path.join(self.image_dir, img_info['file_name'])\n", " image = Image.open(img_path).convert('RGB')\n", " \n", " # Get annotations\n", " anns = self.img_id_to_anns.get(img_id, [])\n", " \n", " # Convert COCO bbox [x, y, w, h] → normalized [cx, cy, w, h]\n", " boxes = []\n", " class_labels = []\n", " for ann in anns:\n", " x, y, w, h = ann['bbox']\n", " # Normalize and convert to center format\n", " cx = (x + w / 2) / img_w\n", " cy = (y + h / 2) / img_h\n", " nw = w / img_w\n", " nh = h / img_h\n", " # Clamp to [0, 1]\n", " cx = max(0.0, min(1.0, cx))\n", " cy = max(0.0, min(1.0, cy))\n", " nw = max(0.001, min(1.0, nw))\n", " nh = max(0.001, min(1.0, nh))\n", " boxes.append([cx, cy, nw, nh])\n", " class_labels.append(ann['category_id'] - 1) # 0-indexed\n", " \n", " # Process image\n", " encoding = self.processor(images=image, return_tensors='pt')\n", " pixel_values = encoding['pixel_values'].squeeze(0) # [3, H, W]\n", " \n", " # Build target\n", " if len(boxes) > 0:\n", " target = {\n", " 'class_labels': torch.tensor(class_labels, dtype=torch.long),\n", " 'boxes': torch.tensor(boxes, dtype=torch.float32),\n", " }\n", " else:\n", " target = {\n", " 'class_labels': torch.zeros(0, dtype=torch.long),\n", " 'boxes': torch.zeros((0, 4), dtype=torch.float32),\n", " }\n", " \n", " return pixel_values, target\n", "\n", "\n", "def collate_fn(batch):\n", " \"\"\"Custom collate: stack pixel_values, keep targets as list.\"\"\"\n", " pixel_values = torch.stack([item[0] for item in batch])\n", " targets = [item[1] for item in batch]\n", " return {'pixel_values': pixel_values}, targets\n", "\n", "\n", "# ── Create datasets and dataloaders ──────────────────────────────────────────\n", "MODEL_PATH = '/mnt/mydata/huggingface/rf-detr-base'\n", "processor = AutoImageProcessor.from_pretrained(MODEL_PATH)\n", "\n", "train_dataset = RobotDetectionDataset(\n", " image_dir=f'{DATASET_DIR}/train',\n", " annotation_file=f'{DATASET_DIR}/train_annotations.json',\n", " processor=processor,\n", " augment=True\n", ")\n", "\n", "val_dataset = RobotDetectionDataset(\n", " image_dir=f'{DATASET_DIR}/val',\n", " annotation_file=f'{DATASET_DIR}/val_annotations.json',\n", " processor=processor,\n", " augment=False\n", ")\n", "\n", "train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True,\n", " collate_fn=collate_fn, num_workers=0)\n", "val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False,\n", " collate_fn=collate_fn, num_workers=0)\n", "\n", "print(f'Train dataset: {len(train_dataset)} images')\n", "print(f'Val dataset: {len(val_dataset)} images')\n", "print(f'Train batches: {len(train_loader)}')\n", "\n", "# Test one batch\n", "batch_inputs, batch_targets = next(iter(train_loader))\n", "print(f'\\nSample batch:')\n", "print(f' pixel_values: {batch_inputs[\"pixel_values\"].shape}')\n", "print(f' targets[0]: class_labels={batch_targets[0][\"class_labels\"]}, '\n", " f'boxes shape={batch_targets[0][\"boxes\"].shape}')" ] }, { "cell_type": "markdown", "id": "model_config_md", "metadata": {}, "source": [ "## 3. Configure Model for Custom Classes\n", "\n", "We need to modify the classification head to output `num_custom_classes + 1` logits\n", "(the +1 is for the background/no-object class).\n", "\n", "**Strategy options:**\n", "1. **Full fine-tune**: Update all weights (best accuracy, needs more data)\n", "2. **Frozen backbone**: Only train decoder + head (faster, good for small datasets)\n", "3. **Partial freeze**: Freeze early backbone layers, fine-tune later layers" ] }, { "cell_type": "code", "execution_count": null, "id": "model_config", "metadata": {}, "outputs": [], "source": [ "from transformers import RfDetrForObjectDetection, RfDetrConfig\n", "import copy\n", "\n", "NUM_CLASSES = len(CLASSES) # 3 custom classes\n", "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n", "\n", "# ── Option A: Load pretrained and replace head ────────────────────────────────\n", "print('Loading pretrained RF-DETR-base...')\n", "model = RfDetrForObjectDetection.from_pretrained(\n", " MODEL_PATH,\n", " num_labels=NUM_CLASSES,\n", " ignore_mismatched_sizes=True # allows head size mismatch\n", ")\n", "\n", "# Update id2label and label2id\n", "model.config.id2label = {i: cls for i, cls in enumerate(CLASSES)}\n", "model.config.label2id = {cls: i for i, cls in enumerate(CLASSES)}\n", "model.config.num_labels = NUM_CLASSES\n", "\n", "model = model.to(device)\n", "\n", "print(f'✅ Model configured for {NUM_CLASSES} classes: {CLASSES}')\n", "print(f' Class head output: {NUM_CLASSES + 1} (classes + background)')\n", "\n", "# ── Freezing strategy ─────────────────────────────────────────────────────────\n", "FREEZE_BACKBONE = True # Set False for full fine-tune\n", "\n", "if FREEZE_BACKBONE:\n", " # Freeze backbone\n", " for param in model.model.backbone.parameters():\n", " param.requires_grad = False\n", " print('\\nFreezing backbone (DINOv2-Small)')\n", "else:\n", " print('\\nFull fine-tune (all parameters trainable)')\n", "\n", "# Count trainable parameters\n", "total = sum(p.numel() for p in model.parameters())\n", "trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n", "frozen = total - trainable\n", "print(f'\\nParameter summary:')\n", "print(f' Total: {total/1e6:.1f}M')\n", "print(f' Trainable: {trainable/1e6:.1f}M ({100*trainable/total:.1f}%)')\n", "print(f' Frozen: {frozen/1e6:.1f}M ({100*frozen/total:.1f}%)')" ] }, { "cell_type": "markdown", "id": "training_md", "metadata": {}, "source": [ "## 4. Training Loop\n", "\n", "RF-DETR uses the **Hungarian matching loss** computed inside the model's forward pass.\n", "When you pass `labels` to the model, it automatically:\n", "1. Runs Hungarian matching between predictions and ground truth\n", "2. Computes classification loss (focal loss)\n", "3. Computes bbox L1 loss\n", "4. Computes GIoU loss\n", "5. Adds auxiliary losses from intermediate decoder layers\n", "\n", "The total loss is returned as `outputs.loss`." ] }, { "cell_type": "code", "execution_count": null, "id": "training_loop", "metadata": {}, "outputs": [], "source": [ "from torch.optim import AdamW\n", "from torch.optim.lr_scheduler import CosineAnnealingLR\n", "import time\n", "\n", "# ── Optimizer & scheduler ─────────────────────────────────────────────────────\n", "# Use different learning rates for backbone vs decoder\n", "backbone_params = [p for n, p in model.named_parameters()\n", " if 'backbone' in n and p.requires_grad]\n", "other_params = [p for n, p in model.named_parameters()\n", " if 'backbone' not in n and p.requires_grad]\n", "\n", "optimizer = AdamW([\n", " {'params': backbone_params, 'lr': 1e-5}, # lower LR for backbone\n", " {'params': other_params, 'lr': 1e-4}, # higher LR for decoder/head\n", "], weight_decay=1e-4)\n", "\n", "N_EPOCHS = 5 # Small for demo; use 50-100 for real training\n", "scheduler = CosineAnnealingLR(optimizer, T_max=N_EPOCHS, eta_min=1e-6)\n", "\n", "print(f'Optimizer: AdamW')\n", "print(f' Backbone LR: 1e-5 ({len(backbone_params)} param groups)')\n", "print(f' Decoder LR: 1e-4 ({len(other_params)} param groups)')\n", "print(f'Scheduler: CosineAnnealingLR (T_max={N_EPOCHS})')\n", "print(f'Epochs: {N_EPOCHS}')\n", "\n", "# ── Training loop ─────────────────────────────────────────────────────────────\n", "train_losses = []\n", "val_losses = []\n", "\n", "for epoch in range(N_EPOCHS):\n", " model.train()\n", " epoch_loss = 0.0\n", " n_batches = 0\n", " t_epoch = time.time()\n", " \n", " for batch_inputs, batch_targets in train_loader:\n", " # Move to device\n", " pixel_values = batch_inputs['pixel_values'].to(device)\n", " labels = [{k: v.to(device) for k, v in t.items()} for t in batch_targets]\n", " \n", " # Forward pass (loss computed internally via Hungarian matching)\n", " outputs = model(pixel_values=pixel_values, labels=labels)\n", " loss = outputs.loss\n", " \n", " # Backward pass\n", " optimizer.zero_grad()\n", " loss.backward()\n", " \n", " # Gradient clipping (important for DETR training stability)\n", " torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.1)\n", " \n", " optimizer.step()\n", " epoch_loss += loss.item()\n", " n_batches += 1\n", " \n", " scheduler.step()\n", " avg_train_loss = epoch_loss / n_batches\n", " train_losses.append(avg_train_loss)\n", " \n", " # Validation\n", " model.eval()\n", " val_loss = 0.0\n", " n_val_batches = 0\n", " with torch.no_grad():\n", " for batch_inputs, batch_targets in val_loader:\n", " pixel_values = batch_inputs['pixel_values'].to(device)\n", " labels = [{k: v.to(device) for k, v in t.items()} for t in batch_targets]\n", " outputs = model(pixel_values=pixel_values, labels=labels)\n", " val_loss += outputs.loss.item()\n", " n_val_batches += 1\n", " \n", " avg_val_loss = val_loss / n_val_batches\n", " val_losses.append(avg_val_loss)\n", " \n", " elapsed = time.time() - t_epoch\n", " lr_now = scheduler.get_last_lr()[1] # decoder LR\n", " print(f'Epoch {epoch+1:2d}/{N_EPOCHS} | '\n", " f'Train Loss: {avg_train_loss:.4f} | '\n", " f'Val Loss: {avg_val_loss:.4f} | '\n", " f'LR: {lr_now:.2e} | '\n", " f'Time: {elapsed:.1f}s')\n", "\n", "print('\\n✅ Training complete!')" ] }, { "cell_type": "code", "execution_count": null, "id": "loss_viz", "metadata": {}, "outputs": [], "source": [ "# ── Plot training curves ──────────────────────────────────────────────────────\n", "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", "fig.patch.set_facecolor('#0d1117')\n", "\n", "epochs_range = range(1, N_EPOCHS + 1)\n", "\n", "ax = axes[0]\n", "ax.set_facecolor('#1a1a2e')\n", "ax.plot(epochs_range, train_losses, 'o-', color='#3498db', linewidth=2,\n", " markersize=6, label='Train Loss')\n", "ax.plot(epochs_range, val_losses, 's-', color='#e74c3c', linewidth=2,\n", " markersize=6, label='Val Loss')\n", "ax.set_title('Training & Validation Loss', color='white', fontweight='bold', fontsize=12)\n", "ax.set_xlabel('Epoch', color='white')\n", "ax.set_ylabel('Loss', color='white')\n", "ax.tick_params(colors='white')\n", "for spine in ax.spines.values(): spine.set_color('white')\n", "ax.legend(facecolor='#1a1a2e', edgecolor='white', labelcolor='white')\n", "ax.grid(True, alpha=0.2, color='white')\n", "\n", "ax2 = axes[1]\n", "ax2.set_facecolor('#1a1a2e')\n", "improvement = [(train_losses[0] - l) / train_losses[0] * 100 for l in train_losses]\n", "ax2.bar(epochs_range, improvement, color='#27ae60', edgecolor='white', linewidth=0.5)\n", "ax2.set_title('Training Loss Improvement (%)', color='white', fontweight='bold', fontsize=12)\n", "ax2.set_xlabel('Epoch', color='white')\n", "ax2.set_ylabel('Improvement vs Epoch 1 (%)', color='white')\n", "ax2.tick_params(colors='white')\n", "for spine in ax2.spines.values(): spine.set_color('white')\n", "ax2.grid(True, alpha=0.2, color='white')\n", "\n", "plt.tight_layout()\n", "plt.savefig('rfdetr_training_curves.png', dpi=120, bbox_inches='tight', facecolor='#0d1117')\n", "plt.show()\n", "print('Training curves saved.')" ] }, { "cell_type": "markdown", "id": "eval_md", "metadata": {}, "source": [ "## 5. Evaluation: Inference on Validation Set" ] }, { "cell_type": "code", "execution_count": null, "id": "eval_code", "metadata": {}, "outputs": [], "source": [ "# ── Run inference on validation images ───────────────────────────────────────\n", "model.eval()\n", "\n", "# Load a few val images for visualization\n", "with open(f'{DATASET_DIR}/val_annotations.json') as f:\n", " val_coco = json.load(f)\n", "\n", "val_img_id_to_anns = {}\n", "for ann in val_coco['annotations']:\n", " val_img_id_to_anns.setdefault(ann['image_id'], []).append(ann)\n", "\n", "fig, axes = plt.subplots(2, 4, figsize=(20, 10))\n", "fig.patch.set_facecolor('#0d1117')\n", "fig.suptitle('RF-DETR Fine-tuned: Validation Set Predictions vs Ground Truth',\n", " color='white', fontsize=13, fontweight='bold')\n", "\n", "for ax_pair_idx, img_info in enumerate(val_coco['images'][:4]):\n", " img = Image.open(f'{DATASET_DIR}/val/{img_info[\"file_name\"]}').convert('RGB')\n", " img_np = np.array(img)\n", " img_h, img_w = img_np.shape[:2]\n", " \n", " # Ground truth\n", " ax_gt = axes[0, ax_pair_idx]\n", " ax_gt.set_facecolor('#0d1117')\n", " ax_gt.imshow(img_np)\n", " ax_gt.set_title(f'GT: img {img_info[\"id\"]}', color='#2ecc71', fontsize=9, fontweight='bold')\n", " ax_gt.axis('off')\n", " for ann in val_img_id_to_anns.get(img_info['id'], []):\n", " x, y, w, h = ann['bbox']\n", " cls_name = CLASSES[ann['category_id'] - 1]\n", " color = list(cat_colors.values())[ann['category_id'] - 1]\n", " rect = patches.Rectangle((x, y), w, h, linewidth=2,\n", " edgecolor=color, facecolor='none')\n", " ax_gt.add_patch(rect)\n", " ax_gt.text(x, y - 3, cls_name, color='white', fontsize=7,\n", " bbox=dict(facecolor=color, alpha=0.8, pad=1, edgecolor='none'))\n", " \n", " # Prediction\n", " ax_pred = axes[1, ax_pair_idx]\n", " ax_pred.set_facecolor('#0d1117')\n", " ax_pred.imshow(img_np)\n", " ax_pred.set_title(f'Pred: img {img_info[\"id\"]}', color='#e74c3c', fontsize=9, fontweight='bold')\n", " ax_pred.axis('off')\n", " \n", " inputs_val = processor(images=img, return_tensors='pt')\n", " inputs_val = {k: v.to(device) for k, v in inputs_val.items()}\n", " with torch.no_grad():\n", " outputs_val = model(**inputs_val)\n", " \n", " target_sizes = torch.tensor([[img_h, img_w]]).to(device)\n", " results = processor.post_process_object_detection(\n", " outputs_val, target_sizes=target_sizes, threshold=0.3\n", " )[0]\n", " \n", " for score, label, box in zip(results['scores'].cpu(),\n", " results['labels'].cpu(),\n", " results['boxes'].cpu()):\n", " x1, y1, x2, y2 = box.tolist()\n", " cls_id = label.item()\n", " cls_name = CLASSES[cls_id] if cls_id < len(CLASSES) else f'cls_{cls_id}'\n", " color = list(cat_colors.values())[cls_id % len(cat_colors)]\n", " rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=2,\n", " edgecolor=color, facecolor='none')\n", " ax_pred.add_patch(rect)\n", " ax_pred.text(x1, y1 - 3, f'{cls_name} {score:.2f}',\n", " color='white', fontsize=7,\n", " bbox=dict(facecolor=color, alpha=0.8, pad=1, edgecolor='none'))\n", "\n", "plt.tight_layout()\n", "plt.savefig('rfdetr_finetuned_predictions.png', dpi=100, bbox_inches='tight', facecolor='#0d1117')\n", "plt.show()\n", "print('Fine-tuned predictions saved.')" ] }, { "cell_type": "markdown", "id": "save_md", "metadata": {}, "source": [ "## 6. Save & Reload Fine-tuned Model" ] }, { "cell_type": "code", "execution_count": null, "id": "save_model", "metadata": {}, "outputs": [], "source": [ "# ── Save fine-tuned model ─────────────────────────────────────────────────────\n", "SAVE_PATH = '/tmp/rfdetr_robot_finetuned'\n", "os.makedirs(SAVE_PATH, exist_ok=True)\n", "\n", "model.save_pretrained(SAVE_PATH)\n", "processor.save_pretrained(SAVE_PATH)\n", "\n", "print(f'✅ Model saved to {SAVE_PATH}')\n", "print('Files saved:')\n", "for fname in sorted(os.listdir(SAVE_PATH)):\n", " size_mb = os.path.getsize(os.path.join(SAVE_PATH, fname)) / 1e6\n", " print(f' {fname:<35} {size_mb:.1f} MB')\n", "\n", "# ── Reload and verify ─────────────────────────────────────────────────────────\n", "print('\\nReloading model...')\n", "model_reloaded = RfDetrForObjectDetection.from_pretrained(SAVE_PATH)\n", "model_reloaded = model_reloaded.to(device)\n", "model_reloaded.eval()\n", "\n", "print(f'✅ Reloaded model classes: {list(model_reloaded.config.id2label.values())}')\n", "\n", "# Quick inference test\n", "img_test = Image.open(f'{DATASET_DIR}/val/{val_coco[\"images\"][0][\"file_name\"]}').convert('RGB')\n", "inputs_test = processor(images=img_test, return_tensors='pt')\n", "inputs_test = {k: v.to(device) for k, v in inputs_test.items()}\n", "with torch.no_grad():\n", " out_test = model_reloaded(**inputs_test)\n", "print(f'✅ Inference test passed: logits shape = {out_test.logits.shape}')" ] }, { "cell_type": "markdown", "id": "tips_md", "metadata": {}, "source": [ "## 7. Training Tips & Best Practices\n", "\n", "### 7.1 Learning Rate Schedule\n", "\n", "```\n", "Backbone LR: 1e-5 (10× lower than decoder — pretrained features)\n", "Decoder LR: 1e-4 (standard for transformer decoder)\n", "Warmup: 500 steps (linear warmup)\n", "Schedule: Cosine annealing or step decay at epoch 40, 60\n", "```\n", "\n", "### 7.2 Data Augmentation\n", "\n", "RF-DETR benefits from strong augmentation:\n", "- **Random horizontal flip** (p=0.5)\n", "- **Random scale** (0.5–2.0×)\n", "- **Random crop** (with box adjustment)\n", "- **Color jitter** (brightness, contrast, saturation)\n", "- **Random erasing** (p=0.3)\n", "\n", "Use `albumentations` or `torchvision.transforms.v2` for bbox-aware augmentation.\n", "\n", "### 7.3 Batch Size & Gradient Accumulation\n", "\n", "```python\n", "# If GPU memory is limited, use gradient accumulation:\n", "ACCUMULATION_STEPS = 4 # effective batch size = batch_size × 4\n", "\n", "for step, (inputs, targets) in enumerate(train_loader):\n", " outputs = model(**inputs, labels=targets)\n", " loss = outputs.loss / ACCUMULATION_STEPS\n", " loss.backward()\n", " \n", " if (step + 1) % ACCUMULATION_STEPS == 0:\n", " torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)\n", " optimizer.step()\n", " optimizer.zero_grad()\n", "```\n", "\n", "### 7.4 Recommended Epochs\n", "\n", "| Dataset size | Epochs | Notes |\n", "|-------------|--------|-------|\n", "| < 100 images | 100-200 | Frozen backbone, strong augmentation |\n", "| 100-1000 | 50-100 | Partial freeze, moderate augmentation |\n", "| > 1000 | 30-50 | Full fine-tune |\n", "| COCO-scale | 12-24 | Standard DETR schedule |\n", "\n", "### 7.5 Gradient Clipping\n", "\n", "**Always use gradient clipping** with DETR-family models:\n", "```python\n", "torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.1)\n", "```\n", "Without this, training can diverge due to large gradients in the attention layers.\n", "\n", "### Next: Notebook 16e — Advanced Features & Visualization\n", "\n", "In the final notebook, we'll explore:\n", "- Attention map visualization inside RF-DETR\n", "- Decoder query analysis\n", "- Comparing DINOv2 features vs RF-DETR features\n", "- Exporting to ONNX for deployment" ] } ] }