""" Simple Training Script for Pest and Disease Classification Using Rich for progress display """ import torch import torch.nn as nn import torch.optim as optim from pathlib import Path import json import argparse from rich.console import Console from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, TimeRemainingColumn from rich.table import Table from rich.panel import Panel from dataset import get_dataloaders, calculate_class_weights from model import create_model console = Console() def train_epoch(model, dataloader, criterion, optimizer, device, progress, task): """Train for one epoch with progress bar""" model.train() running_loss = 0.0 running_corrects = 0 total_samples = 0 for inputs, labels in dataloader: inputs = inputs.to(device) labels = labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) _, preds = torch.max(outputs, 1) loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) total_samples += inputs.size(0) progress.update(task, advance=1) epoch_loss = running_loss / total_samples epoch_acc = running_corrects.double() / total_samples return epoch_loss, epoch_acc.item() def validate_epoch(model, dataloader, criterion, device, progress, task): """Validate for one epoch with progress bar""" model.eval() running_loss = 0.0 running_corrects = 0 total_samples = 0 with torch.no_grad(): for inputs, labels in dataloader: inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) _, preds = torch.max(outputs, 1) running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) total_samples += inputs.size(0) progress.update(task, advance=1) epoch_loss = running_loss / total_samples epoch_acc = running_corrects.double() / total_samples return epoch_loss, epoch_acc.item() def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs, device, save_dir): """ Simple training loop with Rich progress display """ save_dir = Path(save_dir) save_dir.mkdir(exist_ok=True) best_val_acc = 0.0 history = { 'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': [] } console.print("\n[bold green]Starting Training[/bold green]") for epoch in range(num_epochs): console.print(f"\n[bold cyan]Epoch {epoch+1}/{num_epochs}[/bold cyan]") with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), BarColumn(), TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), TimeRemainingColumn(), console=console ) as progress: # Training train_task = progress.add_task( "[red]Training...", total=len(train_loader) ) train_loss, train_acc = train_epoch( model, train_loader, criterion, optimizer, device, progress, train_task ) # Validation val_task = progress.add_task( "[green]Validating...", total=len(val_loader) ) val_loss, val_acc = validate_epoch( model, val_loader, criterion, device, progress, val_task ) # Create results table table = Table(show_header=True, header_style="bold magenta") table.add_column("Split", style="cyan") table.add_column("Loss", justify="right", style="yellow") table.add_column("Accuracy", justify="right", style="green") table.add_row("Train", f"{train_loss:.4f}", f"{train_acc:.4f}") table.add_row("Val", f"{val_loss:.4f}", f"{val_acc:.4f}") console.print(table) # Save history history['train_loss'].append(train_loss) history['train_acc'].append(train_acc) history['val_loss'].append(val_loss) history['val_acc'].append(val_acc) # Save best model if val_acc > best_val_acc: best_val_acc = val_acc torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'val_acc': val_acc, 'val_loss': val_loss, }, save_dir / 'best_model.pth') console.print(f"[bold green]✓ Saved best model (Val Acc: {val_acc:.4f})[/bold green]") # Save checkpoint every 10 epochs if (epoch + 1) % 10 == 0: torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'val_acc': val_acc, 'val_loss': val_loss, }, save_dir / f'checkpoint_epoch_{epoch+1}.pth') console.print(f"[yellow]Checkpoint saved at epoch {epoch+1}[/yellow]") # Save training history with open(save_dir / 'training_history.json', 'w') as f: json.dump(history, f, indent=2) console.print(f"\n[bold green]Training Complete![/bold green]") console.print(f"[bold]Best Val Acc: {best_val_acc:.4f}[/bold]") console.print(f"[bold]Results saved to: {save_dir}/[/bold]") return model, history def main(args): """Main training function""" # Print configuration config_panel = Panel.fit( f"""[bold]Configuration[/bold] Backbone: {args.backbone} Batch Size: {args.batch_size} Image Size: {args.img_size} Epochs: {args.epochs} Learning Rate: {args.lr} Optimizer: {args.optimizer} Device: {args.device} Class Weights: {args.use_class_weights}""", title="Training Settings", border_style="blue" ) console.print(config_panel) # Set device device = torch.device(args.device if torch.cuda.is_available() else 'cpu') console.print(f"\n[bold]Using device: {device}[/bold]") # Load data console.print("\n[bold]Loading datasets...[/bold]") loaders = get_dataloaders( csv_file=args.csv_file, label_mapping_file=args.label_mapping, batch_size=args.batch_size, img_size=args.img_size, num_workers=args.num_workers ) # Create model console.print(f"\n[bold]Creating model: {args.backbone}[/bold]") model = create_model( num_classes=loaders['num_classes'], backbone=args.backbone, pretrained=True, dropout=args.dropout ) model = model.to(device) # Loss function if args.use_class_weights: class_weights = calculate_class_weights(args.csv_file, args.label_mapping) class_weights = class_weights.to(device) criterion = nn.CrossEntropyLoss(weight=class_weights) console.print("[bold]Using weighted CrossEntropyLoss[/bold]") else: criterion = nn.CrossEntropyLoss() console.print("[bold]Using CrossEntropyLoss[/bold]") # Optimizer if args.optimizer == 'adam': optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'adamw': optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'sgd': optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weight_decay) # Train model model, history = train_model( model=model, train_loader=loaders['train'], val_loader=loaders['val'], criterion=criterion, optimizer=optimizer, num_epochs=args.epochs, device=device, save_dir=args.save_dir ) if __name__ == "__main__": parser = argparse.ArgumentParser(description='Simple Training for Pest and Disease Classifier') # Data parameters parser.add_argument('--csv_file', type=str, default='dataset.csv') parser.add_argument('--label_mapping', type=str, default='label_mapping.json') # Model parameters parser.add_argument('--backbone', type=str, default='resnet50', choices=['resnet50', 'resnet101', 'efficientnet_b0', 'efficientnet_b3', 'mobilenet_v2']) parser.add_argument('--dropout', type=float, default=0.3) # Training parameters parser.add_argument('--batch_size', type=int, default=64) parser.add_argument('--img_size', type=int, default=224) parser.add_argument('--epochs', type=int, default=50) parser.add_argument('--lr', type=float, default=0.001) parser.add_argument('--optimizer', type=str, default='adamw', choices=['adam', 'adamw', 'sgd']) parser.add_argument('--weight_decay', type=float, default=0.01) parser.add_argument('--use_class_weights', action='store_true') # System parameters parser.add_argument('--device', type=str, default='cuda', choices=['cuda', 'cpu']) parser.add_argument('--num_workers', type=int, default=8) parser.add_argument('--save_dir', type=str, default='checkpoints') args = parser.parse_args() main(args)