train.py

import os
import sys
import random
import numpy as np
import torch
from torch.utils.data import DataLoader
from configs.default_config import get_arg_parser
from models.segmentor import build_model
from models.losses.ecm_loss import ECMLoss
from models.losses.dice_loss import DiceLoss
from models.losses.combined_loss import CombinedLoss, CrossEntropyDiceLoss
from data.dataset import SegmentationDataset, compute_class_counts_fast
from data.transforms import get_train_transforms, get_val_transforms
from data.sampler import RepeatFactorSampler
from engine.trainer import Trainer


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def build_optimizer(model, args):
    if args.optimizer == 'sgd':
        return torch.optim.SGD(
            model.parameters(), lr=args.lr,
            momentum=args.momentum, weight_decay=args.weight_decay
        )
    elif args.optimizer == 'adam':
        return torch.optim.Adam(
            model.parameters(), lr=args.lr, weight_decay=args.weight_decay
        )
    elif args.optimizer == 'adamw':
        return torch.optim.AdamW(
            model.parameters(), lr=args.lr, weight_decay=args.weight_decay
        )
    raise ValueError(f"Unknown optimizer: {args.optimizer}")


def build_scheduler(optimizer, args):
    if args.scheduler == 'cosine':
        return torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=args.max_epochs, eta_min=1e-6
        )
    elif args.scheduler == 'step':
        return torch.optim.lr_scheduler.StepLR(
            optimizer, step_size=30, gamma=0.1
        )
    elif args.scheduler == 'poly':
        return torch.optim.lr_scheduler.PolynomialLR(
            optimizer, total_iters=args.max_epochs, power=0.9
        )
    raise ValueError(f"Unknown scheduler: {args.scheduler}")


def build_criterion(args, class_counts=None):
    num_classes = args.num_classes
    if args.loss_type == 'ecm':
        return ECMLoss(num_classes, class_counts=class_counts)
    elif args.loss_type == 'combined':
        return CombinedLoss(
            num_classes, class_counts=class_counts,
            ecm_weight=args.ecm_weight, dice_weight=args.dice_weight
        )
    elif args.loss_type == 'ce':
        return torch.nn.CrossEntropyLoss(ignore_index=255)
    elif args.loss_type == 'ce_dice':
        return CrossEntropyDiceLoss(num_classes)
    raise ValueError(f"Unknown loss type: {args.loss_type}")


def main():
    parser = get_arg_parser()
    args = parser.parse_args()
    set_seed(args.seed)
    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    train_transform = get_train_transforms(args.image_size)
    val_transform = get_val_transforms(args.image_size)
    train_dataset = SegmentationDataset(
        args.train_image_dir, args.train_mask_dir,
        args.num_classes, transform=train_transform,
        image_ext=args.image_ext, mask_ext=args.mask_ext
    )
    val_dataset = SegmentationDataset(
        args.val_image_dir, args.val_mask_dir,
        args.num_classes, transform=val_transform,
        image_ext=args.image_ext, mask_ext=args.mask_ext
    )
    class_counts = None
    if args.loss_type in ('ecm', 'combined'):
        print('Computing class counts...')
        class_counts = compute_class_counts_fast(
            train_dataset, args.num_classes, max_samples=args.class_count_samples
        )
        print(f'Class counts: {class_counts}')
    train_sampler = None
    shuffle = True
    if args.use_sampler:
        train_sampler = RepeatFactorSampler(train_dataset, args.num_classes)
        shuffle = False
    train_loader = DataLoader(
        train_dataset, batch_size=args.batch_size,
        shuffle=shuffle, sampler=train_sampler,
        num_workers=args.num_workers, pin_memory=True, drop_last=True
    )
    val_loader = DataLoader(
        val_dataset, batch_size=args.batch_size,
        shuffle=False, num_workers=args.num_workers, pin_memory=True
    )
    model = build_model(
        num_classes=args.num_classes,
        backbone_variant=args.backbone_variant,
        fpn_channels=args.fpn_channels,
        seg_inner_channels=args.seg_inner_channels,
        use_light_head=args.use_light_head
    )
    model = model.to(device)
    criterion = build_criterion(args, class_counts)
    if hasattr(criterion, 'to'):
        criterion = criterion.to(device)
    optimizer = build_optimizer(model, args)
    scheduler = build_scheduler(optimizer, args)
    trainer = Trainer(
        model=model, criterion=criterion, optimizer=optimizer,
        scheduler=scheduler, device=device,
        train_loader=train_loader, val_loader=val_loader,
        num_classes=args.num_classes, max_epochs=args.max_epochs,
        log_dir=args.log_dir, save_dir=args.save_dir,
        save_interval=args.save_interval, eval_interval=args.eval_interval,
        grad_clip=args.grad_clip
    )
    if args.resume is not None:
        trainer.resume(args.resume)
    trainer.train()


if __name__ == '__main__':
    main()