Create train.py

train.py

@@ -0,0 +1,239 @@
+# ================================================================
+#  🔄  IMAGE ROTATION PREDICTION — From-Scratch ResNet-18
+#      Dataset: ImageNet-1k  ·  Hardware: Kaggle T4 GPU
+# ================================================================
+
+!pip install -q transformers datasets
+
+# ────────────────────── Imports ──────────────────────
+import os, random, math, time
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from transformers import ResNetConfig, ResNetForImageClassification
+from datasets import load_dataset
+from tqdm.auto import tqdm
+
+# ────────────────────── Config ───────────────────────
+HF_TOKEN        = "hf_XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
+NUM_TRAIN       = 50_000
+NUM_VAL         = 5_000
+IMG_SIZE        = 224
+BATCH_SIZE      = 128
+EPOCHS          = 12
+LR              = 1e-3
+WARMUP_EPOCHS   = 1
+WEIGHT_DECAY    = 0.05
+LABEL_SMOOTHING = 0.1
+DEVICE          = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+TRAIN_DIR       = "/kaggle/working/data/train"
+VAL_DIR         = "/kaggle/working/data/val"
+MODEL_DIR       = "/kaggle/working/rotation_model"
+
+print(f"🖥️  Device: {DEVICE}")
+if DEVICE.type == "cuda":
+    print(f"    GPU:  {torch.cuda.get_device_name()}")
+    print(f"    VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")
+
+# ──────────── Download ImageNet-1k (Streaming) ──────────────
+from huggingface_hub import login
+login(token=HF_TOKEN)
+
+def download_images(split, save_dir, num_images):
+    os.makedirs(save_dir, exist_ok=True)
+    existing = len([f for f in os.listdir(save_dir) if f.endswith(".jpg")])
+    if existing >= num_images:
+        print(f"  ✓ {save_dir}: {existing} images already exist → skipping.")
+        return
+    ds = load_dataset("ILSVRC/imagenet-1k", split=split,
+                      streaming=True, trust_remote_code=True, token=HF_TOKEN)
+    count = 0
+    for ex in tqdm(ds, total=num_images, desc=f"  ↓ {split}"):
+        if count >= num_images:
+            break
+        try:
+            img = ex["image"].convert("RGB")
+            w, h = img.size
+            if min(w, h) > 480:
+                s = 480 / min(w, h)
+                img = img.resize((int(w*s), int(h*s)), Image.BILINEAR)
+            img.save(os.path.join(save_dir, f"{count}.jpg"), quality=90)
+            count += 1
+        except Exception:
+            continue
+    print(f"  ✓ {count} Images → {save_dir}")
+
+print("\n📥 Loading images from ImageNet-1k …")
+download_images("train",      TRAIN_DIR, NUM_TRAIN)
+download_images("validation", VAL_DIR,   NUM_VAL)
+
+# ──────────────────── Rotation-Dataset ───────────────────────
+ANGLES = [0, 90, 180, 270]
+ANGLE_NAMES = ["0° (original)", "90° CCW", "180°", "270° CCW (=90° CW)"]
+
+class RotationDataset(Dataset):
+    def __init__(self, img_dir, num_imgs, transform, all_rotations=False):
+        self.img_dir  = img_dir
+        self.num_imgs = num_imgs
+        self.transform = transform
+        self.all_rot  = all_rotations
+
+    def __len__(self):
+        return self.num_imgs * 4 if self.all_rot else self.num_imgs
+
+    def __getitem__(self, idx):
+        if self.all_rot:
+            img_idx, label = idx // 4, idx % 4
+        else:
+            img_idx, label = idx, random.randint(0, 3)
+
+        img = Image.open(os.path.join(self.img_dir, f"{img_idx}.jpg")).convert("RGB")
+
+        angle = ANGLES[label]
+        if   angle == 90:  img = img.transpose(Image.ROTATE_90)
+        elif angle == 180: img = img.transpose(Image.ROTATE_180)
+        elif angle == 270: img = img.transpose(Image.ROTATE_270)
+
+        return self.transform(img), label
+
+train_tf = transforms.Compose([
+    transforms.Resize(256),
+    transforms.RandomCrop(IMG_SIZE),
+    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.05),
+    transforms.RandomGrayscale(p=0.05),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    transforms.RandomErasing(p=0.1),
+])
+val_tf = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(IMG_SIZE),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+train_ds = RotationDataset(TRAIN_DIR, NUM_TRAIN, train_tf, all_rotations=True)
+val_ds   = RotationDataset(VAL_DIR,   NUM_VAL,   val_tf,   all_rotations=True)
+
+train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,
+                          num_workers=2, pin_memory=True, drop_last=True)
+val_loader   = DataLoader(val_ds,   batch_size=BATCH_SIZE, shuffle=False,
+                          num_workers=2, pin_memory=True)
+
+print(f"\n📊 Dataset size:")
+print(f"   Train: {len(train_ds):>8,}  ({NUM_TRAIN:,} images × 4 rotations)")
+print(f"   Val:   {len(val_ds):>8,}  ({NUM_VAL:,} images × 4 rotations)")
+
+# ────────────────── Modell: ResNet-18 from scratch ───────────────
+config = ResNetConfig(
+    num_channels=3,
+    embedding_size=64,
+    hidden_sizes=[64, 128, 256, 512],   # 4 Stages
+    depths=[2, 2, 2, 2],                # → ResNet-18
+    layer_type="basic",
+    hidden_act="relu",
+    num_labels=4,                        # 0°, 90°, 180°, 270°
+)
+model = ResNetForImageClassification(config).to(DEVICE)
+n_params = sum(p.numel() for p in model.parameters())
+print(f"\n🏗️  Model: ResNet-18 from scratch — {n_params:,} parameters")
+
+# ────────────────────── Training-Setup ───────────────────────
+optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
+
+total_steps  = len(train_loader) * EPOCHS
+warmup_steps = len(train_loader) * WARMUP_EPOCHS
+
+def lr_lambda(step):
+    if step < warmup_steps:
+        return step / max(warmup_steps, 1)
+    progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
+    return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+scaler    = torch.cuda.amp.GradScaler()
+criterion = nn.CrossEntropyLoss(label_smoothing=LABEL_SMOOTHING)
+
+# ────────────────────── Training-Loop ────────────────────────
+best_val_acc = 0.0
+print(f"\n🚀 Starting training: {EPOCHS} epochs, {total_steps:,} steps\n")
+
+for epoch in range(EPOCHS):
+    t0 = time.time()
+
+    # ---- Train ----
+    model.train()
+    run_loss = correct = total = 0
+
+    pbar = tqdm(train_loader, desc=f"Ep {epoch+1:2d}/{EPOCHS} [Train]", leave=False)
+    for imgs, labels in pbar:
+        imgs   = imgs.to(DEVICE, non_blocking=True)
+        labels = labels.to(DEVICE, non_blocking=True)
+
+        with torch.cuda.amp.autocast():
+            logits = model(pixel_values=imgs).logits
+            loss   = criterion(logits, labels)
+
+        optimizer.zero_grad(set_to_none=True)
+        scaler.scale(loss).backward()
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        scaler.step(optimizer)
+        scaler.update()
+        scheduler.step()
+
+        bs = labels.size(0)
+        run_loss += loss.item() * bs
+        correct  += (logits.argmax(1) == labels).sum().item()
+        total    += bs
+        pbar.set_postfix(loss=f"{run_loss/total:.4f}", acc=f"{100*correct/total:.1f}%")
+
+    train_acc = 100 * correct / total
+
+    # ---- Validate ----
+    model.eval()
+    v_correct = v_total = 0
+    v_loss = 0.0
+    cls_correct = [0]*4
+    cls_total   = [0]*4
+
+    with torch.no_grad():
+        for imgs, labels in tqdm(val_loader, desc=f"Ep {epoch+1:2d}/{EPOCHS} [Val]  ", leave=False):
+            imgs   = imgs.to(DEVICE, non_blocking=True)
+            labels = labels.to(DEVICE, non_blocking=True)
+            with torch.cuda.amp.autocast():
+                logits = model(pixel_values=imgs).logits
+                loss   = criterion(logits, labels)
+            preds = logits.argmax(1)
+            bs = labels.size(0)
+            v_loss    += loss.item() * bs
+            v_correct += (preds == labels).sum().item()
+            v_total   += bs
+            for c in range(4):
+                mask = (labels == c)
+                cls_correct[c] += (preds[mask] == labels[mask]).sum().item()
+                cls_total[c]   += mask.sum().item()
+
+    val_acc = 100 * v_correct / v_total
+    dt = time.time() - t0
+
+    print(f"Epoch {epoch+1:2d}/{EPOCHS} │ "
+          f"Train {train_acc:.1f}% │ Val {val_acc:.2f}% │ "
+          f"LR {scheduler.get_last_lr()[0]:.6f} │ {dt:.0f}s")
+    for c in range(4):
+        ca = 100*cls_correct[c]/max(cls_total[c],1)
+        print(f"   {ANGLE_NAMES[c]:>25s}: {ca:.1f}%")
+
+    if val_acc > best_val_acc:
+        best_val_acc = val_acc
+        model.save_pretrained(MODEL_DIR)
+        print(f"   ✅ New best model saved → {MODEL_DIR}")
+    print()
+
+# ── Fertig ──
+print("=" * 60)
+print(f"🏆 Training finished!  Best Val-Accuracy: {best_val_acc:.2f}%")
+print(f"   Model: {MODEL_DIR}")