# scripts/train.py """ Script di training per il classificatore basilico vs pomodoro. FunzionalitΓ : - 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 sys import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights from torch.utils.data import DataLoader from pathlib import Path # 1) Percorsi dataset (usa percorsi assoluti per sicurezza) script_dir = Path(__file__).parent base_dir = script_dir.parent if script_dir.parent.name != "scripts" else script_dir train_dir = base_dir / "scripts" / "data" / "basil_tomato" / "train" val_dir = base_dir / "scripts" / "data" / "basil_tomato" / "val" models_dir = base_dir / "scripts" / "models" print(f"πŸ” Cercando dataset in:") print(f" Train: {train_dir}") print(f" Val: {val_dir}") print(f" Models: {models_dir}") # Verifica esistenza directory if not train_dir.exists(): print(f"❌ Directory train non trovata: {train_dir}") sys.exit(1) if not val_dir.exists(): print(f"❌ Directory validation non trovata: {val_dir}") sys.exit(1) # 2) Valori standard di normalizzazione ImageNet IMGNET_MEAN = [0.485, 0.456, 0.406] IMGNET_STD = [0.229, 0.224, 0.225] # 3) Trasformazioni dati train_transforms = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(IMGNET_MEAN, IMGNET_STD) ]) val_transforms = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(IMGNET_MEAN, IMGNET_STD) ]) # 4) Dataset e DataLoader con error handling try: train_ds = datasets.ImageFolder(str(train_dir), transform=train_transforms) val_ds = datasets.ImageFolder(str(val_dir), transform=val_transforms) if len(train_ds) == 0: print(f"❌ Nessuna immagine trovata in {train_dir}") sys.exit(1) if len(val_ds) == 0: print(f"❌ Nessuna immagine trovata in {val_dir}") sys.exit(1) except Exception as e: print(f"❌ Errore nel caricamento dataset: {e}") sys.exit(1) # Ottimizza batch size per GPU disponibile device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if torch.cuda.is_available(): gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3 batch_size = 32 if gpu_memory > 6 else 16 num_workers = min(4, os.cpu_count() or 1) print(f"πŸš€ GPU: {torch.cuda.get_device_name(0)} ({gpu_memory:.1f} GB)") print(f"βš™οΈ Batch size ottimizzato: {batch_size}") else: batch_size = 8 num_workers = min(2, os.cpu_count() or 1) print("πŸ’» Usando CPU") train_loader = DataLoader( train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=torch.cuda.is_available() ) val_loader = DataLoader( val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=torch.cuda.is_available() ) print(f"βœ… Classi trovate: {train_ds.classes}") print(f"πŸ“Š Numero immagini - Train: {len(train_ds)}, Validation: {len(val_ds)}") # Verifica bilanciamento classi class_counts_train = {} class_counts_val = {} for idx, (_, label) in enumerate(train_ds): class_name = train_ds.classes[label] class_counts_train[class_name] = class_counts_train.get(class_name, 0) + 1 for idx, (_, label) in enumerate(val_ds): class_name = val_ds.classes[label] class_counts_val[class_name] = class_counts_val.get(class_name, 0) + 1 print(f"πŸ“ˆ Distribuzione train: {class_counts_train}") print(f"πŸ“ˆ Distribuzione val: {class_counts_val}") # 5) Configura device (giΓ  fatto sopra) # 6) Costruisci il modello con pesi pre-addestrati (fix deprecation warning) print("πŸ”„ Caricando EfficientNet-B0 con pesi pre-addestrati...") try: model = efficientnet_b0(weights=EfficientNet_B0_Weights.IMAGENET1K_V1) num_classes = len(train_ds.classes) # Sostituisci il classificatore finale model.classifier[1] = nn.Linear( model.classifier[1].in_features, num_classes ) model = model.to(device) # Ottimizzazioni per GPU if torch.cuda.is_available(): model = model.half() # Usa mixed precision per risparmiare memoria print("βœ… Mixed precision attivata") print(f"βœ… Modello caricato con {num_classes} classi: {train_ds.classes}") except Exception as e: print(f"❌ Errore nel caricamento del modello: {e}") sys.exit(1) # 7) Criterio e ottimizzatore criterion = nn.CrossEntropyLoss() optimizer = optim.Adam( model.parameters(), lr=1e-4, weight_decay=1e-5 ) # 8) Funzione di training per un'epoca con progress tracking def train_epoch(): model.train() running_loss, running_corrects = 0.0, 0 total_batches = len(train_loader) for batch_idx, (inputs, labels) in enumerate(train_loader): inputs, labels = inputs.to(device), labels.to(device) # Mixed precision per GPU if torch.cuda.is_available(): inputs = inputs.half() 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() # Progress tracking if (batch_idx + 1) % max(1, total_batches // 10) == 0: progress = (batch_idx + 1) / total_batches * 100 print(f" πŸ“ˆ Training progress: {progress:.1f}% ({batch_idx + 1}/{total_batches})") epoch_loss = running_loss / len(train_ds) epoch_acc = running_corrects / len(train_ds) return epoch_loss, epoch_acc # 9) Funzione di validazione con progress tracking def validate_epoch(): model.eval() val_loss, val_corrects = 0.0, 0 total_batches = len(val_loader) with torch.no_grad(): for batch_idx, (inputs, labels) in enumerate(val_loader): inputs, labels = inputs.to(device), labels.to(device) # Mixed precision per GPU if torch.cuda.is_available(): inputs = inputs.half() outputs = model(inputs) loss = criterion(outputs, labels) val_loss += loss.item() * inputs.size(0) val_corrects += (outputs.argmax(1) == labels).sum().item() # Progress tracking if (batch_idx + 1) % max(1, total_batches // 5) == 0: progress = (batch_idx + 1) / total_batches * 100 print(f" πŸ“Š Validation progress: {progress:.1f}% ({batch_idx + 1}/{total_batches})") loss = val_loss / len(val_ds) acc = val_corrects / len(val_ds) return loss, acc # 10) Loop di training principale con miglioramenti if __name__ == "__main__": import time best_val_acc = 0.0 models_dir.mkdir(exist_ok=True) print(f"\nπŸš€ Iniziando training per {10} epoche...") print(f"πŸ’Ύ I modelli saranno salvati in: {models_dir}") start_time = time.time() for epoch in range(1, 11): # 10 epoche epoch_start = time.time() print(f"\nπŸ”„ Epoca {epoch}/10:") # Training print(" πŸ‹οΈ Training...") train_loss, train_acc = train_epoch() # Validation print(" πŸ” Validation...") val_loss, val_acc = validate_epoch() epoch_time = time.time() - epoch_start print( f"βœ… Epoca {epoch}: train_loss={train_loss:.4f}, " f"train_acc={train_acc:.4f} | val_loss={val_loss:.4f}, " f"val_acc={val_acc:.4f} | tempo={epoch_time:.1f}s" ) # Salva il miglior modello con validazione if val_acc > best_val_acc: best_val_acc = val_acc save_path = models_dir / "basil_tomato_classifier.pth" try: # Salva sia state_dict che modello completo torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'best_val_acc': best_val_acc, 'train_acc': train_acc, 'val_loss': val_loss, 'classes': train_ds.classes, 'num_classes': num_classes }, save_path) print(f"πŸ’Ύ Nuovo best model salvato con val_acc={val_acc:.4f}") except Exception as e: print(f"❌ Errore nel salvataggio: {e}") # Cleanup GPU memory if torch.cuda.is_available(): torch.cuda.empty_cache() total_time = time.time() - start_time print(f"\nπŸŽ‰ Training completato! Best val_acc: {best_val_acc:.4f}") print(f"⏱️ Tempo totale: {total_time:.1f}s ({total_time/60:.1f} minuti)") # Statistiche finali if torch.cuda.is_available(): print(f"πŸ“Š Memoria GPU utilizzata: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB") torch.cuda.reset_peak_memory_stats()