base code test for training a model with basil and tomatoes images to include in step 3 of previous code to improve answer

test2_with_training/scripts/train.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+# scripts/train.py
+"""
+Script di training per il classificatore basilico vs pomodoro.
+Struttura:
+  - carica dataset da data/basil_tomato/train e /val
+  - transfer learning con EfficientNet-B0
+  - salva il miglior modello in models/basil_tomato_classifier.pth
+"""
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms, models
+from torch.utils.data import DataLoader
+
+# 1) Percorsi dataset
+train_dir = "data/basil_tomato/train"
+val_dir   = "data/basil_tomato/val"
+
+# 2) Trasformazioni dati
+train_transforms = transforms.Compose([
+    transforms.RandomResizedCrop(224),
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+val_transforms = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+# 3) Crea dataset e DataLoader
+train_ds = datasets.ImageFolder(train_dir, transform=train_transforms)
+val_ds   = datasets.ImageFolder(val_dir,   transform=val_transforms)
+
+train_loader = DataLoader(train_ds, batch_size=32, shuffle=True,  num_workers=4)
+val_loader   = DataLoader(val_ds,   batch_size=32, shuffle=False, num_workers=4)
+
+print(f"Classi trovate: {train_ds.classes}")
+print(f"Numero immagini train: {len(train_ds)}, validation: {len(val_ds)}")
+
+# 4) Configura device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# 5) Costruisci il modello
+model = models.efficientnet_b0(pretrained=True)
+num_classes = len(train_ds.classes)
+model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
+model = model.to(device)
+
+# 6) Definisci criterio e ottimizzatore
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
+
+# 7) Funzioni di training e validation
+
+def train_epoch():
+    model.train()
+    running_loss, running_corrects = 0.0, 0
+
+    for inputs, labels in train_loader:
+        inputs, labels = inputs.to(device), labels.to(device)
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        running_loss     += loss.item() * inputs.size(0)
+        running_corrects += (outputs.argmax(1) == labels).sum().item()
+
+    epoch_loss = running_loss / len(train_ds)
+    epoch_acc  = running_corrects / len(train_ds)
+    return epoch_loss, epoch_acc
+
+
+def validate_epoch():
+    model.eval()
+    val_loss, val_corrects = 0.0, 0
+
+    with torch.no_grad():
+        for inputs, labels in val_loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+
+            val_loss     += loss.item() * inputs.size(0)
+            val_corrects += (outputs.argmax(1) == labels).sum().item()
+
+    loss = val_loss / len(val_ds)
+    acc  = val_corrects / len(val_ds)
+    return loss, acc
+
+# 8) Training loop principale
+best_val_acc = 0.0
+os.makedirs("models", exist_ok=True)
+
+for epoch in range(1, 11):  # 10 epoche
+    train_loss, train_acc = train_epoch()
+    val_loss,   val_acc   = validate_epoch()
+
+    print(f"Epoca {epoch}: train_loss={train_loss:.4f}, train_acc={train_acc:.4f} | "
+          f"val_loss={val_loss:.4f}, val_acc={val_acc:.4f}")
+
+    # Salva il modello migliore
+    if val_acc > best_val_acc:
+        best_val_acc = val_acc
+        save_path = os.path.join("models", "basil_tomato_classifier.pth")
+        torch.save(model.state_dict(), save_path)
+        print(f"--> Nuovo best model salvato con val_acc={val_acc:.4f}")
+
+print("Training completato. Best val_acc: {:.4f}".format(best_val_acc))