| import os | |
| import json | |
| import torch | |
| import wandb | |
| import datetime | |
| import numpy as np | |
| from tqdm import tqdm | |
| from torch.utils.data import DataLoader | |
| from torch.utils.tensorboard import SummaryWriter | |
| from segmentation_models_pytorch.base.modules import Activation | |
| from SemanticModel.data_loader import SegmentationDataset | |
| from SemanticModel.metrics import compute_mean_iou | |
| from SemanticModel.image_preprocessing import get_training_augmentations, get_validation_augmentations | |
| from SemanticModel.utilities import list_images, validate_dimensions | |
| class ModelTrainer: | |
| def __init__(self, model_config, root_dir, epochs=40, train_size=1024, | |
| val_size=None, workers=2, batch_size=2, learning_rate=1e-4, | |
| step_count=2, decay_factor=0.8, wandb_config=None, | |
| optimizer='rmsprop', target_class=None, resume_path=None): | |
| self.config = model_config | |
| self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| self.root_dir = root_dir | |
| self._initialize_training_params(epochs, train_size, val_size, workers, | |
| batch_size, learning_rate, step_count, | |
| decay_factor, optimizer, target_class) | |
| self._setup_directories() | |
| self._initialize_datasets() | |
| self._setup_optimizer() | |
| self._initialize_tracking() | |
| if resume_path: | |
| self._resume_training(resume_path) | |
| def _initialize_training_params(self, epochs, train_size, val_size, workers, | |
| batch_size, learning_rate, step_count, | |
| decay_factor, optimizer, target_class): | |
| self.epochs = epochs | |
| self.train_size = train_size | |
| self.val_size = val_size | |
| self.workers = workers | |
| self.batch_size = batch_size | |
| self.learning_rate = learning_rate | |
| self.step_schedule = self._calculate_step_schedule(epochs, step_count) | |
| self.decay_factor = decay_factor | |
| self.optimizer_type = optimizer | |
| self.target_class = target_class | |
| self.current_epoch = 1 | |
| self.best_iou = 0.0 | |
| self.best_epoch = 0 | |
| self.classes = ['background'] + self.config.classes if self.config.background_flag else self.config.classes | |
| def _setup_directories(self): | |
| """Verifies and creates necessary directories.""" | |
| self.train_dir = os.path.join(self.root_dir, 'train') | |
| self.val_dir = os.path.join(self.root_dir, 'val') | |
| required_subdirs = ['Images', 'Masks'] | |
| for path in [self.train_dir] + ([self.val_dir] if os.path.exists(self.val_dir) else []): | |
| for subdir in required_subdirs: | |
| full_path = os.path.join(path, subdir) | |
| if not os.path.exists(full_path): | |
| raise FileNotFoundError(f"Missing directory: {full_path}") | |
| def _initialize_datasets(self): | |
| """Sets up training and validation datasets.""" | |
| self.train_dataset = SegmentationDataset( | |
| self.train_dir, | |
| classes=self.classes, | |
| augmentation=get_training_augmentations(self.train_size, self.train_size), | |
| preprocessing=self.config.preprocessing | |
| ) | |
| if os.path.exists(self.val_dir): | |
| self.val_dataset = SegmentationDataset( | |
| self.val_dir, | |
| classes=self.classes, | |
| augmentation=get_validation_augmentations( | |
| self.val_size or self.train_size, | |
| self.val_size or self.train_size, | |
| fixed_size=False | |
| ), | |
| preprocessing=self.config.preprocessing | |
| ) | |
| self.val_loader = DataLoader( | |
| self.val_dataset, | |
| batch_size=1, | |
| shuffle=False, | |
| num_workers=self.workers | |
| ) | |
| else: | |
| self.val_dataset = self.train_dataset | |
| self.val_loader = DataLoader( | |
| self.val_dataset, | |
| batch_size=1, | |
| shuffle=False, | |
| num_workers=self.workers | |
| ) | |
| self.train_loader = DataLoader( | |
| self.train_dataset, | |
| batch_size=self.batch_size, | |
| shuffle=True, | |
| num_workers=self.workers | |
| ) | |
| def _setup_optimizer(self): | |
| """Configures model optimizer.""" | |
| optimizer_map = { | |
| 'adam': torch.optim.Adam, | |
| 'sgd': lambda params: torch.optim.SGD(params, momentum=0.9), | |
| 'rmsprop': torch.optim.RMSprop | |
| } | |
| optimizer_class = optimizer_map.get(self.optimizer_type.lower()) | |
| if not optimizer_class: | |
| raise ValueError(f"Unsupported optimizer: {self.optimizer_type}") | |
| self.optimizer = optimizer_class([{'params': self.config.model.parameters(), | |
| 'lr': self.learning_rate}]) | |
| def _initialize_tracking(self): | |
| """Sets up training progress tracking.""" | |
| timestamp = datetime.datetime.now().strftime("%m-%d-%Y_%H%M%S") | |
| self.output_dir = os.path.join( | |
| self.root_dir, | |
| f'model_outputs-{self.config.architecture}[{self.config.encoder}]-{timestamp}' | |
| ) | |
| os.makedirs(self.output_dir, exist_ok=True) | |
| self.writer = SummaryWriter(log_dir=self.output_dir) | |
| self.metrics = { | |
| 'best_epoch': self.best_epoch, | |
| 'best_epoch_iou': self.best_iou, | |
| 'last_epoch': 0, | |
| 'last_epoch_iou': 0.0, | |
| 'last_epoch_lr': self.learning_rate, | |
| 'step_schedule': self.step_schedule, | |
| 'decay_factor': self.decay_factor, | |
| 'target_class': self.target_class or 'overall' | |
| } | |
| def _calculate_step_schedule(self, epochs, steps): | |
| """Calculates learning rate step schedule.""" | |
| return list(map(int, np.linspace(0, epochs, steps + 2)[1:-1])) | |
| def train(self): | |
| """Executes training loop.""" | |
| model = self.config.model.to(self.device) | |
| if torch.cuda.device_count() > 1: | |
| model = torch.nn.DataParallel(model) | |
| print(f'Using {torch.cuda.device_count()} GPUs') | |
| self._save_config() | |
| for epoch in range(self.current_epoch, self.epochs + 1): | |
| print(f'\nEpoch {epoch}/{self.epochs}') | |
| print(f'Learning rate: {self.optimizer.param_groups[0]["lr"]:.3e}') | |
| train_loss = self._train_epoch(model) | |
| val_loss, val_metrics = self._validate_epoch(model) | |
| self._update_tracking(epoch, train_loss, val_loss, val_metrics) | |
| self._adjust_learning_rate(epoch) | |
| self._save_checkpoints(model, epoch, val_metrics) | |
| print(f'\nTraining completed. Best {self.metrics["target_class"]} IoU: {self.best_iou:.3f}') | |
| return model, self.metrics | |
| def _train_epoch(self, model): | |
| """Executes single training epoch.""" | |
| model.train() | |
| total_loss = 0 | |
| sample_count = 0 | |
| for batch in tqdm(self.train_loader, desc='Training'): | |
| images, masks = [x.to(self.device) for x in batch] | |
| self.optimizer.zero_grad() | |
| outputs = model(images) | |
| loss = self.config.loss(outputs, masks) | |
| loss.backward() | |
| self.optimizer.step() | |
| total_loss += loss.item() * len(images) | |
| sample_count += len(images) | |
| return total_loss / sample_count | |
| def _validate_epoch(self, model): | |
| """Executes validation pass.""" | |
| model.eval() | |
| total_loss = 0 | |
| predictions = [] | |
| ground_truth = [] | |
| with torch.no_grad(): | |
| for batch in tqdm(self.val_loader, desc='Validation'): | |
| images, masks = [x.to(self.device) for x in batch] | |
| outputs = model(images) | |
| loss = self.config.loss(outputs, masks) | |
| total_loss += loss.item() | |
| if self.config.n_classes > 1: | |
| predictions.extend([p.cpu().argmax(dim=0) for p in outputs]) | |
| ground_truth.extend([m.cpu().argmax(dim=0) for m in masks]) | |
| else: | |
| predictions.extend([(torch.sigmoid(p) > 0.5).float().squeeze().cpu() | |
| for p in outputs]) | |
| ground_truth.extend([m.cpu().squeeze() for m in masks]) | |
| metrics = compute_mean_iou( | |
| predictions, | |
| ground_truth, | |
| num_classes=len(self.classes), | |
| ignore_index=255 | |
| ) | |
| return total_loss / len(self.val_loader), metrics | |
| def _update_tracking(self, epoch, train_loss, val_loss, val_metrics): | |
| """Updates training metrics and logging.""" | |
| mean_iou = val_metrics['mean_iou'] | |
| print(f"\nLosses - Train: {train_loss:.3f}, Val: {val_loss:.3f}") | |
| print(f"Mean IoU: {mean_iou:.3f}") | |
| self.writer.add_scalar('Loss/train', train_loss, epoch) | |
| self.writer.add_scalar('Loss/val', val_loss, epoch) | |
| self.writer.add_scalar('IoU/mean', mean_iou, epoch) | |
| for idx, iou in enumerate(val_metrics['per_category_iou']): | |
| print(f"{self.classes[idx]} IoU: {iou:.3f}") | |
| self.writer.add_scalar(f'IoU/{self.classes[idx]}', iou, epoch) | |
| def _adjust_learning_rate(self, epoch): | |
| """Adjusts learning rate according to schedule.""" | |
| if epoch in self.step_schedule: | |
| current_lr = self.optimizer.param_groups[0]['lr'] | |
| new_lr = current_lr * self.decay_factor | |
| for param_group in self.optimizer.param_groups: | |
| param_group['lr'] = new_lr | |
| print(f'\nDecreased learning rate: {current_lr:.3e} -> {new_lr:.3e}') | |
| def _save_checkpoints(self, model, epoch, metrics): | |
| """Saves model checkpoints and metrics.""" | |
| epoch_iou = (metrics['mean_iou'] if self.target_class is None | |
| else metrics['per_category_iou'][self.classes.index(self.target_class)]) | |
| self.metrics.update({ | |
| 'last_epoch': epoch, | |
| 'last_epoch_iou': round(float(epoch_iou), 3), | |
| 'last_epoch_lr': self.optimizer.param_groups[0]['lr'] | |
| }) | |
| if epoch_iou > self.best_iou: | |
| self.best_iou = epoch_iou | |
| self.best_epoch = epoch | |
| self.metrics.update({ | |
| 'best_epoch': epoch, | |
| 'best_epoch_iou': round(float(epoch_iou), 3), | |
| 'overall_iou': round(float(metrics['mean_iou']), 3) | |
| }) | |
| torch.save(model, os.path.join(self.output_dir, 'best_model.pth')) | |
| print(f'New best model saved (IoU: {epoch_iou:.3f})') | |
| torch.save(model, os.path.join(self.output_dir, 'last_model.pth')) | |
| with open(os.path.join(self.output_dir, 'metrics.json'), 'w') as f: | |
| json.dump(self.metrics, f, indent=4) | |
| def _save_config(self): | |
| """Saves training configuration.""" | |
| config = { | |
| **self.config.config_data, | |
| 'train_size': self.train_size, | |
| 'val_size': self.val_size, | |
| 'epochs': self.epochs, | |
| 'batch_size': self.batch_size, | |
| 'optimizer': self.optimizer_type, | |
| 'workers': self.workers, | |
| 'target_class': self.target_class or 'overall' | |
| } | |
| with open(os.path.join(self.output_dir, 'config.json'), 'w') as f: | |
| json.dump(config, f, indent=4) | |
| def _resume_training(self, resume_path): | |
| """Resumes training from checkpoint.""" | |
| if not os.path.exists(resume_path): | |
| raise FileNotFoundError(f"Resume path not found: {resume_path}") | |
| required_files = { | |
| 'model': 'last_model.pth', | |
| 'metrics': 'metrics.json', | |
| 'config': 'config.json' | |
| } | |
| paths = {k: os.path.join(resume_path, v) for k, v in required_files.items()} | |
| if not all(os.path.exists(p) for p in paths.values()): | |
| raise FileNotFoundError("Missing required checkpoint files") | |
| with open(paths['config']) as f: | |
| config = json.load(f) | |
| with open(paths['metrics']) as f: | |
| metrics = json.load(f) | |
| self.current_epoch = metrics['last_epoch'] + 1 | |
| self.best_iou = metrics['best_epoch_iou'] | |
| self.best_epoch = metrics['best_epoch'] | |
| self.learning_rate = metrics['last_epoch_lr'] | |
| print(f'Resuming training from epoch {self.current_epoch}') |