""" AgriFM training script for PASTIS - v2. Fixes: augmentation, class weights, small model, stronger regularization. AMD ROCm compatible. Usage: python train.py --small_model --fold 1 --epochs 100 python train.py --tiny_model --fold 1 --epochs 100 python train.py --fold 1 --epochs 100 (base 196M) """ import os import sys import time import argparse import json import math from pathlib import Path from datetime import datetime, timedelta import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader from torch.amp import GradScaler, autocast torch.serialization.add_safe_globals([np.core.multiarray.scalar]) sys.path.insert(0, str(Path(__file__).parent)) from models.agrifm import (build_agrifm_pastis, build_agrifm_pastis_small, build_agrifm_pastis_tiny) from datasets.pastis_dataset import (PASTISDataset, NUM_CLASSES, IGNORE_INDEX, PASTIS_CLASSES) from losses.loss import CropCELoss, compute_class_weights from evaluation.metrics import SegmentationMetrics # --------------------------------------------------------------------------- # Display helpers # --------------------------------------------------------------------------- CYAN = '\033[96m' GREEN = '\033[92m' YELLOW = '\033[93m' RED = '\033[91m' BOLD = '\033[1m' RESET = '\033[0m' LINE = '─' * 70 DLINE = '═' * 70 def header(text): print(f"\n{BOLD}{CYAN}{DLINE}{RESET}") print(f"{BOLD}{CYAN} {text}{RESET}") print(f"{BOLD}{CYAN}{DLINE}{RESET}") def section(text): print(f"\n{BOLD}{YELLOW}{LINE}{RESET}") print(f"{BOLD}{YELLOW} {text}{RESET}") print(f"{BOLD}{YELLOW}{LINE}{RESET}") def fmt_time(seconds): return str(timedelta(seconds=int(seconds))) def progress_bar(current, total, width=40): filled = int(width * current / total) bar = '█' * filled + '░' * (width - filled) return f"[{bar}] {100*current/total:.1f}%" # --------------------------------------------------------------------------- # Logger — writes plain text to file, color to console # --------------------------------------------------------------------------- class Logger: def __init__(self, log_path): self.log_path = log_path with open(log_path, 'w') as f: f.write(f"AgriFM PASTIS Training Log\n") f.write(f"Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") f.write("=" * 70 + "\n\n") def write(self, text): clean = text for code in [CYAN, GREEN, YELLOW, RED, BOLD, RESET]: clean = clean.replace(code, '') with open(self.log_path, 'a') as f: f.write(clean + '\n') def log_args(self, args): self.write("TRAINING CONFIGURATION") self.write("=" * 40) for k, v in vars(args).items(): self.write(f" {k:<24} = {v}") self.write("") def log_epoch(self, epoch, train_loss, val_metrics, epoch_time, lr, is_best): self.write(f"Epoch {epoch:03d} [{fmt_time(epoch_time)}] lr={lr:.2e}") self.write(f" Train Loss : {train_loss:.4f}") if val_metrics: for k in ['OA','mIoU','mFscore','mPrecision','mRecall','Kappa']: self.write(f" {k:<14}: {val_metrics.get(k, 0.):.2f}%") self.write(f" Val Loss : {val_metrics.get('val_loss', 0.):.4f}") if is_best: self.write( f" *** NEW BEST mFscore: " f"{val_metrics.get('mFscore', 0.):.2f}% ***" ) self.write("") def log_final(self, test_metrics, best_mfscore, total_time): self.write("=" * 70) self.write("FINAL TEST RESULTS") self.write("=" * 70) for k in ['OA','mIoU','mFscore','mPrecision','mRecall','Kappa']: self.write(f" {k:<14}: {test_metrics.get(k, 0.):.2f}%") self.write(f" Test Loss : {test_metrics.get('val_loss', 0.):.4f}") self.write(f" Best Val mFscore : {best_mfscore:.2f}%") self.write(f" Total Time : {fmt_time(total_time)}") self.write( f" Finished: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" ) self.write("=" * 70) # --------------------------------------------------------------------------- # Arguments # --------------------------------------------------------------------------- def parse_args(): p = argparse.ArgumentParser('AgriFM PASTIS Training') # Data p.add_argument('--data_root', default='/workspace/project/PASTIS') p.add_argument('--fold', type=int, default=1, choices=[1,2,3,4,5]) p.add_argument('--num_frames', type=int, default=32) p.add_argument('--label_channel', type=int, default=0, help='0=coarse(20cls), 1=medium(43cls)') p.add_argument('--num_classes', type=int, default=20) # Model size — mutually exclusive mg = p.add_mutually_exclusive_group() mg.add_argument('--small_model', action='store_true', help='~14M params — recommended for PASTIS from scratch') mg.add_argument('--tiny_model', action='store_true', help='~5M params — use if small still overfits') p.add_argument('--use_checkpoint', action='store_true', help='Gradient checkpointing — saves VRAM') p.add_argument('--pretrained', default=None, help='Path to pretrained AgriFM weights') # Training p.add_argument('--epochs', type=int, default=100) p.add_argument('--batch_size', type=int, default=16) p.add_argument('--lr', type=float, default=5e-5) p.add_argument('--weight_decay', type=float, default=0.05) p.add_argument('--warmup_iters', type=int, default=500) p.add_argument('--amp', action='store_true', default=True) p.add_argument('--no_amp', dest='amp', action='store_false') p.add_argument('--no_class_weights', action='store_true', help='Disable class weighting') # System p.add_argument('--num_workers', type=int, default=4) p.add_argument('--work_dir', default='./work_dirs/pastis_fold1') p.add_argument('--resume', default=None) p.add_argument('--log_interval', type=int, default=50) p.add_argument('--val_interval', type=int, default=1) return p.parse_args() # --------------------------------------------------------------------------- # LR Scheduler: linear warmup → cosine decay # --------------------------------------------------------------------------- class WarmupCosineScheduler: def __init__(self, optimizer, warmup_iters, total_iters, base_lr, min_lr=1e-6): self.optimizer = optimizer self.warmup_iters = warmup_iters self.total_iters = total_iters self.base_lr = base_lr self.min_lr = min_lr self.current_iter = 0 def step(self): self.current_iter += 1 it = self.current_iter if it <= self.warmup_iters: lr = (self.min_lr + (self.base_lr - self.min_lr) * (it / self.warmup_iters)) else: progress = ((it - self.warmup_iters) / max(1, self.total_iters - self.warmup_iters)) lr = (self.min_lr + 0.5 * (self.base_lr - self.min_lr) * (1 + math.cos(math.pi * progress))) for pg in self.optimizer.param_groups: pg['lr'] = lr return lr def get_lr(self): for pg in self.optimizer.param_groups: return pg['lr'] # --------------------------------------------------------------------------- # Train one epoch # --------------------------------------------------------------------------- def train_epoch(model, loader, optimizer, scheduler, criterion, device, scaler, epoch, args, total_epochs): model.train() total_loss = 0. n_batches = 0 t_start = time.time() t_batch = time.time() print(f"\n {BOLD}Training{RESET}") for i, batch in enumerate(loader): s2 = batch['S2'].to(device, non_blocking=True) label = batch['label'].to(device, non_blocking=True) optimizer.zero_grad(set_to_none=True) with autocast('cuda', enabled=args.amp): logits = model(s2) loss_dict = criterion(logits, label) loss = loss_dict['loss'] scaler.scale(loss).backward() scaler.unscale_(optimizer) nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0) scaler.step(optimizer) scaler.update() lr = scheduler.step() total_loss += loss.item() n_batches += 1 if (i+1) % args.log_interval == 0 or (i+1) == len(loader): batch_time = time.time() - t_batch remaining = len(loader) - (i+1) eta = (batch_time / max(1, args.log_interval)) * remaining avg_loss = total_loss / n_batches print(f" {progress_bar(i+1, len(loader))} " f"iter {i+1:4d}/{len(loader)} " f"loss={avg_loss:.4f} " f"lr={lr:.2e} " f"eta={fmt_time(eta)}") t_batch = time.time() avg_loss = total_loss / max(n_batches, 1) epoch_time = time.time() - t_start print(f" {GREEN}Done: avg_loss={avg_loss:.4f} " f"time={fmt_time(epoch_time)}{RESET}") return avg_loss, epoch_time # --------------------------------------------------------------------------- # Validation / Test # --------------------------------------------------------------------------- @torch.no_grad() def validate(model, loader, criterion, device, num_classes, ignore_index, args, split='Val'): model.eval() metrics = SegmentationMetrics(num_classes, ignore_index) total_loss = 0. n_batches = 0 t0 = time.time() print(f"\n {BOLD}{split}idating...{RESET}") for batch in loader: s2 = batch['S2'].to(device, non_blocking=True) label = batch['label'].to(device, non_blocking=True) with autocast('cuda', enabled=args.amp): logits = model(s2) loss_dict = criterion(logits, label) total_loss += loss_dict['loss'].item() n_batches += 1 metrics.update(logits.float(), label) result = metrics.compute() result['val_loss'] = round(total_loss / max(n_batches, 1), 4) result['val_time'] = round(time.time() - t0, 1) return result, metrics # --------------------------------------------------------------------------- # Print epoch metrics table # --------------------------------------------------------------------------- def print_metrics(train_loss, val_metrics, epoch, total_epochs, total_time, best_mfscore, is_best): section(f"Epoch {epoch+1:03d} / {total_epochs:03d} | " f"Elapsed: {fmt_time(total_time)}") print(f"\n {'Metric':<18} {'Value':>10}") print(f" {'─'*18} {'─'*10}") print(f" {'Train Loss':<18} {train_loss:>9.4f}") print(f" {'Val Loss':<18} {val_metrics.get('val_loss', 0.):>9.4f}") print(f" {'─'*18} {'─'*10}") for key in ['OA','mIoU','mFscore','mPrecision','mRecall','Kappa']: val = val_metrics.get(key, 0.) color = GREEN if key in ['mFscore','mIoU'] else RESET star = ' ★' if (key == 'mFscore' and is_best) else '' print(f" {color}{BOLD}{key:<18}{RESET} " f"{color}{BOLD}{val:>9.2f}%{RESET}{star}") print(f" {'─'*18} {'─'*10}") print(f" {'Best mFscore':<18} {best_mfscore:>9.2f}%") print(f" {'Val Time':<18} {val_metrics.get('val_time', 0.):>9.1f}s") # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): args = parse_args() start_time = time.time() os.makedirs(args.work_dir, exist_ok=True) log_path = os.path.join(args.work_dir, 'training_log.txt') logger = Logger(log_path) # Determine model size label if args.small_model: model_size = 'small (~14M)' elif args.tiny_model: model_size = 'tiny (~5M)' else: model_size = 'base (196M)' header(f"AgriFM × PASTIS — Fold {args.fold} | " f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') gpu_name = (torch.cuda.get_device_name(0) if device.type == 'cuda' else 'CPU') vram_gb = (torch.cuda.get_device_properties(0).total_memory / 1024**3 if device.type == 'cuda' else 0) print(f"\n {'Device':<22}: {gpu_name}") print(f" {'VRAM':<22}: {vram_gb:.1f} GB") print(f" {'Model size':<22}: {model_size}") print(f" {'AMP':<22}: {args.amp}") print(f" {'Batch size':<22}: {args.batch_size}") print(f" {'Epochs':<22}: {args.epochs}") print(f" {'LR':<22}: {args.lr}") print(f" {'Weight decay':<22}: {args.weight_decay}") print(f" {'Warmup iters':<22}: {args.warmup_iters}") print(f" {'Augmentation':<22}: True (flip + rotate)") print(f" {'Class weights':<22}: {not args.no_class_weights}") print(f" {'Work dir':<22}: {args.work_dir}") # Save args with open(os.path.join(args.work_dir, 'args.json'), 'w') as f: json.dump(vars(args), f, indent=2) logger.log_args(args) logger.write(f"Model size: {model_size}") # ── Datasets ──────────────────────────────────────────────────────── section("Building Datasets") train_ds = PASTISDataset( args.data_root, fold=args.fold, split='train', num_frames=args.num_frames, label_channel=args.label_channel, augment=True) val_ds = PASTISDataset( args.data_root, fold=args.fold, split='val', num_frames=args.num_frames, label_channel=args.label_channel, augment=False) test_ds = PASTISDataset( args.data_root, fold=args.fold, split='test', num_frames=args.num_frames, label_channel=args.label_channel, augment=False) train_loader = DataLoader( train_ds, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True, drop_last=True, persistent_workers=(args.num_workers > 0)) val_loader = DataLoader( val_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True, persistent_workers=(args.num_workers > 0)) test_loader = DataLoader( test_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) print(f"\n Train: {len(train_ds)} " f"Val: {len(val_ds)} " f"Test: {len(test_ds)}") print(f" Train batches: {len(train_loader)} " f"Val batches: {len(val_loader)}") # ── Model ─────────────────────────────────────────────────────────── section("Building Model") if args.small_model: model = build_agrifm_pastis_small( num_classes=args.num_classes, use_checkpoint=args.use_checkpoint) elif args.tiny_model: model = build_agrifm_pastis_tiny( num_classes=args.num_classes, use_checkpoint=args.use_checkpoint) else: model = build_agrifm_pastis( num_classes=args.num_classes, use_checkpoint=args.use_checkpoint) if args.pretrained: model.load_pretrained(args.pretrained) model = model.to(device) n_params = sum(p.numel() for p in model.parameters()) / 1e6 n_train = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6 print(f"\n Total params : {n_params:.1f}M") print(f" Trainable params : {n_train:.1f}M") logger.write(f"Model: {n_params:.1f}M params ({model_size})") # ── Class weights ─────────────────────────────────────────────────── cls_weights = None if not args.no_class_weights: section("Computing Class Weights") cls_weights = compute_class_weights( args.data_root, fold=args.fold, label_channel=args.label_channel, num_classes=args.num_classes, ignore_index=IGNORE_INDEX, max_samples=300, ).to(device) # ── Loss ──────────────────────────────────────────────────────────── criterion = CropCELoss( ignore_index=IGNORE_INDEX, hard_mining_ratio=0.5, class_weights=cls_weights, ) # ── Optimizer ─────────────────────────────────────────────────────── optimizer = torch.optim.AdamW( model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.999), ) # ── Scheduler ─────────────────────────────────────────────────────── total_iters = args.epochs * len(train_loader) scheduler = WarmupCosineScheduler( optimizer, warmup_iters=args.warmup_iters, total_iters=total_iters, base_lr=args.lr, min_lr=1e-6, ) # ── AMP scaler ────────────────────────────────────────────────────── scaler = GradScaler('cuda', enabled=args.amp) # ── Resume ────────────────────────────────────────────────────────── start_epoch = 0 best_mfscore = 0. if args.resume: section(f"Resuming from {args.resume}") ckpt = torch.load(args.resume, map_location=device, weights_only=False) model.load_state_dict(ckpt['model']) optimizer.load_state_dict(ckpt['optimizer']) start_epoch = ckpt['epoch'] + 1 best_mfscore = ckpt.get('best_mfscore', 0.) print(f" Resumed epoch {start_epoch}, " f"best mFscore={best_mfscore:.2f}%") all_logs = [] # ── Training loop ──────────────────────────────────────────────────── header(f"Training | {args.epochs} epochs | " f"model={model_size} | fold={args.fold}") for epoch in range(start_epoch, args.epochs): train_loss, epoch_time = train_epoch( model, train_loader, optimizer, scheduler, criterion, device, scaler, epoch, args, args.epochs ) if epoch % args.val_interval == 0: val_metrics, _ = validate( model, val_loader, criterion, device, args.num_classes, IGNORE_INDEX, args, split='Val' ) else: val_metrics = {} total_elapsed = time.time() - start_time mfscore = val_metrics.get('mFscore', 0.) is_best = mfscore > best_mfscore lr_now = scheduler.get_lr() print_metrics(train_loss, val_metrics, epoch, args.epochs, total_elapsed, best_mfscore, is_best) logger.log_epoch(epoch, train_loss, val_metrics, epoch_time, lr_now, is_best) all_logs.append({ 'epoch': epoch, 'train_loss': round(train_loss, 4), 'lr': lr_now, 'epoch_time': round(epoch_time, 1), **val_metrics, }) # Save best checkpoint if is_best: best_mfscore = mfscore torch.save({ 'epoch': epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'best_mfscore': best_mfscore, 'val_metrics': val_metrics, 'args': vars(args), }, os.path.join(args.work_dir, 'best_model.pth')) print(f"\n {GREEN}{BOLD}★ New best mFscore: " f"{best_mfscore:.2f}% → saved best_model.pth{RESET}") # Save latest checkpoint torch.save({ 'epoch': epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'best_mfscore': best_mfscore, 'args': vars(args), }, os.path.join(args.work_dir, 'latest.pth')) # Save per-epoch log with open(os.path.join(args.work_dir, 'log.json'), 'w') as f: json.dump(all_logs, f, indent=2) # ── Final test evaluation ──────────────────────────────────────────── header("Final Test Evaluation") best_ckpt = torch.load( os.path.join(args.work_dir, 'best_model.pth'), map_location=device, weights_only=False ) model.load_state_dict(best_ckpt['model']) test_metrics, test_metric_obj = validate( model, test_loader, criterion, device, args.num_classes, IGNORE_INDEX, args, split='Test' ) total_time = time.time() - start_time per_cls_iou = test_metric_obj.per_class_iou() # Print final results section("Final Test Results") print(f"\n {'Metric':<18} {'Value':>10}") print(f" {'─'*18} {'─'*10}") for key in ['OA','mIoU','mFscore','mPrecision','mRecall','Kappa']: val = test_metrics.get(key, 0.) color = GREEN if key in ['mFscore','mIoU'] else RESET print(f" {color}{BOLD}{key:<18}{RESET} " f"{color}{BOLD}{val:>9.2f}%{RESET}") print(f" {'─'*18} {'─'*10}") print(f" {'Test Loss':<18} " f"{test_metrics.get('val_loss', 0.):>9.4f}") print(f" {'Total Time':<18} {fmt_time(total_time):>10}") # Per-class IoU table section("Per-Class IoU (Test)") print(f"\n {'Class':<30} {'IoU':>7} Bar") print(f" {'─'*30} {'─'*7} {'─'*20}") for c in range(args.num_classes): if c == IGNORE_INDEX: continue name = PASTIS_CLASSES.get(c, f'Class_{c}') iou = per_cls_iou[c] * 100 bar = '█' * int(iou / 5) color = GREEN if iou > 20 else (YELLOW if iou > 5 else RED) print(f" {color}{name:<30} {iou:>7.2f}% {bar}{RESET}") logger.log_final(test_metrics, best_mfscore, total_time) # Save all results results = { 'test_metrics': test_metrics, 'best_val_mfscore': best_mfscore, 'total_time_s': round(total_time, 1), 'model_size': model_size, 'per_class_iou': { PASTIS_CLASSES.get(c, f'Class_{c}'): round(float(per_cls_iou[c] * 100), 2) for c in range(args.num_classes) if c != IGNORE_INDEX }, 'args': vars(args), } with open(os.path.join(args.work_dir, 'test_results.json'), 'w') as f: json.dump(results, f, indent=2) print(f"\n {GREEN}{BOLD}All results saved to: {args.work_dir}{RESET}") print(f" training_log.txt — full training log") print(f" log.json — per-epoch metrics") print(f" test_results.json — final test + per-class IoU") print(f" best_model.pth — best checkpoint (by val mFscore)") print(f" latest.pth — last epoch checkpoint") print(f" args.json — training configuration") if __name__ == '__main__': main()