Upload new.py

new.py

@@ -0,0 +1,148 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+import timm
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader, Subset
+from sklearn.model_selection import train_test_split
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+# =========================
+# CONFIG
+# =========================
+DATA_DIR = r"D:\BPA PROJECT\train"
+BATCH_SIZE = 32  # Tiny model allows for larger batches = faster training
+EPOCHS = 30
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+print(f"Checking hardware... Using: {DEVICE}")
+
+# =========================
+# MODEL
+# =========================
+class BreedClassifier(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+        print("📥 Downloading/Loading ConvNeXt-Tiny weights (this may take a minute)...")
+        # Tiny is much faster than Base but still extremely accurate
+        self.model = timm.create_model(
+            "convnext_tiny.fb_in22k_ft_in1k", 
+            pretrained=True, 
+            num_classes=num_classes
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+# =========================
+# TRANSFORMS
+# =========================
+norm_stats = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+
+train_tfms = transforms.Compose([
+    transforms.RandomResizedCrop(224),
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),
+    transforms.Normalize(*norm_stats)
+])
+
+val_tfms = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(*norm_stats)
+])
+
+# =========================
+# DATA LOADING
+# =========================
+def load_data():
+    print("📂 Scanning dataset folders...")
+    full_dataset = datasets.ImageFolder(DATA_DIR)
+    class_names = full_dataset.classes
+    
+    indices = list(range(len(full_dataset)))
+    train_idx, val_idx = train_test_split(
+        indices, 
+        test_size=0.2, 
+        stratify=full_dataset.targets,
+        random_state=42
+    )
+
+    train_ds = Subset(datasets.ImageFolder(DATA_DIR, transform=train_tfms), train_idx)
+    val_ds = Subset(datasets.ImageFolder(DATA_DIR, transform=val_tfms), val_idx)
+
+    # Set num_workers=0 if you are on Windows to avoid multi-processing hangs
+    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
+    val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, num_workers=0)
+
+    return train_loader, val_loader, len(class_names), class_names
+
+# =========================
+# TRAINING
+# =========================
+def train():
+    train_loader, val_loader, num_classes, class_names = load_data()
+    
+    model = BreedClassifier(num_classes).to(DEVICE)
+
+    # Higher learning rate for the smaller 'tiny' model
+    optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05)
+    criterion = nn.CrossEntropyLoss() 
+    scheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS)
+
+    print(f"✅ Setup complete. Starting {EPOCHS} epochs of training...")
+
+    best_acc = 0
+    for epoch in range(EPOCHS):
+        model.train()
+        total_loss = 0
+        correct = total = 0
+
+        for batch_idx, (x, y) in enumerate(train_loader):
+            x, y = x.to(DEVICE), y.to(DEVICE)
+
+            optimizer.zero_grad()
+            out = model(x)
+            loss = criterion(out, y)
+            loss.backward()
+            optimizer.step()
+
+            pred = out.argmax(1)
+            correct += (pred == y).sum().item()
+            total += y.size(0)
+            total_loss += loss.item()
+            
+            if batch_idx % 5 == 0:
+                print(f"Epoch {epoch+1} | Batch {batch_idx}/{len(train_loader)} | Loss: {loss.item():.4f}", end='\r')
+
+        # VALIDATION
+        model.eval()
+        val_correct = val_total = val_conf = 0
+        with torch.no_grad():
+            for x, y in val_loader:
+                x, y = x.to(DEVICE), y.to(DEVICE)
+                out = model(x)
+                probs = F.softmax(out, dim=1)
+                conf, pred = torch.max(probs, dim=1)
+
+                val_correct += (pred == y).sum().item()
+                val_total += y.size(0)
+                val_conf += conf.sum().item()
+
+        val_acc = 100 * val_correct / val_total
+        avg_conf = 100 * val_conf / val_total
+        
+        print(f"\n✨ Epoch [{epoch+1}/{EPOCHS}] - Val Acc: {val_acc:.2f}% | Confidence: {avg_conf:.2f}%")
+        
+        scheduler.step()
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+            torch.save(model.state_dict(), "best_breed_model.pth")
+            print("💾 Model Saved!")
+
+if __name__ == "__main__":
+    train()