train.py

"""Training entry for the lightweight runtime.

Usage:
    Single GPU:
        python train.py <config.py> --device cuda:0

    Multi GPU on one node:
        torchrun --nproc_per_node=4 train.py <config.py> --launcher pytorch

Notes:
    - Distributed training requires CUDA. The actual device is always
      resolved from `LOCAL_RANK`.
    - `train_dataloader.sampler` should be a distributed-aware sampler
      such as `DefaultSampler`.
    - Logging is handled by `MMLogger`. Rank 0 writes `work_dir/train.log`
      and non-zero ranks only emit `ERROR` logs to stdout by default.
"""

import argparse
import logging
import os
import random
import sys
from datetime import datetime
from math import ceil

import numpy as np
import torch

# Add project root to path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from mmengine.config import Config
from mmengine.logging import MMLogger, MessageHub
from mmengine.runner import load_checkpoint as load_model_checkpoint
from mmengine.utils.dl_utils import set_multi_processing

from runtime.builders import (build_dataloader_from_cfg, build_log_processor,
                              build_optimizer_and_wrapper, build_scheduler,
                              build_visualizer, resolve_train_loop_cfg,
                              validate_scheduler_loop_compatibility)
from runtime.checkpoint import load_scheduler_state_dict
from runtime.common import require_dataset_meta, to_plain_dict
from runtime.context import RuntimeContext
from runtime.dist import (barrier, init_distributed_runtime, resolve_device,
                          sync_string_from_rank0,
                          wrap_model_for_distributed_training)
from runtime.hooks import (build_runtime_hooks,
                           validate_runtime_hook_requirements)
from runtime.loops import EpochBasedTrainLoop, IterBasedTrainLoop, ValLoop
from utils.builder import build_dataset, build_model

INFINITE_TRAIN_SAMPLERS = {
    'InfiniteSampler',
    'MultiSourceSampler',
    'GroupMultiSourceSampler',
}


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description='Train detection model')
    parser.add_argument('config', help='config file path')
    parser.add_argument('--work-dir', help='working directory')
    parser.add_argument('--resume-from', help='checkpoint path to resume from')
    parser.add_argument('--load-from', help='pretrained checkpoint path')
    parser.add_argument(
        '--device',
        default='cuda',
        help='device for single-process training; distributed training '
        'always resolves the concrete CUDA device from LOCAL_RANK')
    parser.add_argument(
        '--launcher',
        choices=('none', 'pytorch', 'slurm', 'mpi'),
        default='none',
        help='job launcher for distributed training')
    parser.add_argument(
        '--local-rank',
        '--local_rank',
        dest='local_rank',
        type=int,
        default=0,
        help=argparse.SUPPRESS)
    return parser.parse_args()


def prepare_work_dir(cfg, config_path: str, work_dir, timestamp: str) -> None:
    config_name = os.path.splitext(os.path.basename(config_path))[0]

    base_work_dir = work_dir or cfg.get('work_dir', './work_dirs')
    if not base_work_dir:
        base_work_dir = './work_dirs'

    # Create per-run directory: <base>/<config>/<timestamp>.
    if os.path.basename(os.path.normpath(base_work_dir)) == config_name:
        cfg.work_dir = os.path.join(base_work_dir, timestamp)
    else:
        cfg.work_dir = os.path.join(base_work_dir, config_name, timestamp)
    os.makedirs(cfg.work_dir, exist_ok=True)


def build_logger(work_dir):
    log_file = os.path.join(work_dir, 'train.log')
    return MMLogger.get_instance(
        name='Detector',
        logger_name='Detector',
        log_file=log_file,
        log_level=logging.INFO,
        distributed=False)


def setup_runtime_environment(cfg, distributed: bool) -> None:
    env_cfg = to_plain_dict(cfg.get('env_cfg', {}))
    if not isinstance(env_cfg, dict):
        raise TypeError(f'env_cfg must be a dict, but got {type(env_cfg)}')

    mp_cfg = env_cfg.get('mp_cfg', {})
    if not isinstance(mp_cfg, dict):
        raise TypeError(f'env_cfg.mp_cfg must be a dict, but got {type(mp_cfg)}')

    set_multi_processing(
        mp_start_method=mp_cfg.get('mp_start_method', 'fork'),
        opencv_num_threads=int(mp_cfg.get('opencv_num_threads', 0)),
        distributed=distributed)
    torch.backends.cudnn.benchmark = bool(
        env_cfg.get('cudnn_benchmark', False))


def setup_randomness(cfg, logger) -> None:
    randomness = to_plain_dict(cfg.get('randomness', {}))
    if not isinstance(randomness, dict):
        raise TypeError(
            f'randomness must be a dict, but got {type(randomness)}')

    seed = randomness.get('seed', None)
    deterministic = bool(randomness.get('deterministic', False))
    if seed is None:
        logger.info('No randomness.seed configured; RNG state is left unchanged.')
    else:
        seed = int(seed)
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
        logger.info(f'Random seed set to {seed}')

    torch.backends.cudnn.deterministic = deterministic
    if deterministic:
        logger.info('CUDNN deterministic mode enabled.')


def log_runtime_start(logger, config_path: str, cfg, distributed: bool) -> None:
    logger.info(f'Starting training, config: {config_path}')
    logger.info(f'Using device: {cfg.device}')
    if distributed:
        logger.info(
            f'Distributed training enabled: rank={os.environ["RANK"]}, '
            f'world_size={os.environ["WORLD_SIZE"]}, '
            f'local_rank={os.environ["LOCAL_RANK"]}')


def resolve_checkpoint_options(args, cfg):
    resume_from = args.resume_from or cfg.get('resume_from', None)
    load_from = args.load_from or cfg.get('load_from', None)
    if bool(cfg.get('resume', False)) and resume_from is None:
        raise NotImplementedError(
            'cfg.resume=True auto-resume is not supported by this runtime. '
            'Set --resume-from or cfg.resume_from explicitly.')
    if resume_from is not None and load_from is not None:
        raise ValueError(
            'resume_from and load_from are mutually exclusive. Use resume_from '
            'to continue training, or load_from to initialize model weights.')
    checkpoint_cfg = to_plain_dict(cfg.get('checkpoint_cfg', {}))
    resume_strict = bool(
        checkpoint_cfg.get('resume_strict', cfg.get('resume_strict', True)))
    load_strict = bool(
        checkpoint_cfg.get('load_strict', cfg.get('load_strict', True)))
    return resume_from, load_from, resume_strict, load_strict


def build_runtime_model(cfg, logger, resume_from, distributed: bool):
    logger.info('Building model...')
    model = build_model(cfg.model)
    model = model.to(cfg.device)

    if resume_from is None:
        if hasattr(model, 'init_weights'):
            logger.info('Initializing model weights...')
            model.init_weights()

    if distributed:
        model = wrap_model_for_distributed_training(model, cfg)

    return model


def _sampler_type(loader_cfg) -> str:
    sampler_cfg = loader_cfg.get('sampler', None)
    if not isinstance(sampler_cfg, dict):
        raise KeyError(
            'dataloader.sampler must be an explicit dict with a type field.')
    sampler_type = sampler_cfg.get('type', None)
    if sampler_type is None:
        raise KeyError('dataloader.sampler.type is required.')
    return str(sampler_type).split('.')[-1]


def validate_dataloader_cfg(loader_cfg, loop_type: str, name: str) -> None:
    sampler_type = _sampler_type(loader_cfg)
    if name == 'train_dataloader' and loop_type == 'EpochBasedTrainLoop':
        if sampler_type in INFINITE_TRAIN_SAMPLERS:
            raise ValueError(
                f'{sampler_type} is an iteration-based sampler and cannot be '
                'used with EpochBasedTrainLoop. Use IterBasedTrainLoop or set '
                'an explicit finite sampler such as DefaultSampler.')
    if name != 'train_dataloader' and sampler_type in INFINITE_TRAIN_SAMPLERS:
        raise ValueError(
            f'{name} must not use iteration-based sampler {sampler_type}.')


def ensure_non_empty_dataloader(dataloader, name: str) -> None:
    num_batches = len(dataloader)
    if num_batches <= 0:
        raise ValueError(
            f'{name} produced zero batches. Check dataset length, batch_size, '
            'drop_last, sampler, and batch_sampler explicitly.')


def build_data_loaders(cfg, loop_type: str, logger):
    logger.info('Building data loaders...')
    train_dl_cfg = to_plain_dict(cfg.get('train_dataloader', None))
    val_dl_cfg = to_plain_dict(cfg.get('val_dataloader', None))
    if train_dl_cfg is None or val_dl_cfg is None:
        raise KeyError('Config must define train_dataloader and val_dataloader.')
    validate_dataloader_cfg(train_dl_cfg, loop_type, 'train_dataloader')
    validate_dataloader_cfg(val_dl_cfg, loop_type, 'val_dataloader')

    train_dataset = build_dataset(train_dl_cfg['dataset'])
    val_dataset = build_dataset(val_dl_cfg['dataset'])

    train_dataloader = build_dataloader_from_cfg(
        train_dl_cfg,
        train_dataset,
        default_shuffle=True)
    val_dataloader = build_dataloader_from_cfg(
        val_dl_cfg, val_dataset, default_shuffle=False)
    ensure_non_empty_dataloader(train_dataloader, 'train_dataloader')
    ensure_non_empty_dataloader(val_dataloader, 'val_dataloader')

    return train_dataset, val_dataset, train_dataloader, val_dataloader


def build_runtime_context(cfg,
                          logger,
                          model,
                          train_dataset,
                          val_dataset,
                          train_dataloader,
                          val_dataloader,
                          train_loop_cfg,
                          timestamp: str) -> RuntimeContext:
    loop_type = train_loop_cfg['type']
    val_interval = int(train_loop_cfg.get('val_interval', 1))
    if val_interval <= 0:
        raise ValueError('train_cfg.val_interval must be > 0')
    val_metric_key = cfg.get('val_metric_key',
                             train_loop_cfg.get('val_metric_key', None))
    if val_metric_key is not None and not isinstance(val_metric_key, str):
        raise TypeError('val_metric_key must be a string when configured.')

    visualizer_dataset_meta = require_dataset_meta(
        val_dataset, val_dataloader, 'val_dataloader.dataset')
    visualizer = build_visualizer(
        cfg, work_dir=cfg.work_dir, dataset_meta=visualizer_dataset_meta)

    optim_wrapper = build_optimizer_and_wrapper(model, cfg, logger)
    param_schedulers = build_scheduler(
        cfg, optim_wrapper=optim_wrapper, iters_per_epoch=len(train_dataloader))
    validate_scheduler_loop_compatibility(loop_type, param_schedulers)
    if param_schedulers:
        scheduler_names = ', '.join(
            type(scheduler).__name__ for scheduler in param_schedulers)
        logger.info(f'Built param schedulers: {scheduler_names}')

    log_processor = build_log_processor(cfg)

    message_hub = MessageHub.get_instance('fundet_train')
    message_hub.update_info('epoch', 0)
    message_hub.update_info('iter', 0)

    if loop_type == 'IterBasedTrainLoop':
        max_iters = int(train_loop_cfg['max_iters'])
        max_epochs = ceil(max_iters / len(train_dataloader))
    else:
        max_epochs = int(train_loop_cfg['max_epochs'])
        max_iters = max_epochs * len(train_dataloader)

    hook_manager = build_runtime_hooks(cfg, logger)
    context = RuntimeContext(
        model=model,
        cfg=cfg,
        work_dir=cfg.work_dir,
        logger=logger,
        message_hub=message_hub,
        optim_wrapper=optim_wrapper,
        param_schedulers=param_schedulers,
        log_processor=log_processor,
        visualizer=visualizer,
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        val_dataset=val_dataset,
        hook_manager=hook_manager,
        loop_type=loop_type,
        val_interval=val_interval,
        val_metric_key=val_metric_key,
        max_epochs=max_epochs,
        max_iters=max_iters,
        timestamp=timestamp)
    validate_runtime_hook_requirements(hook_manager, loop_type)
    return context


def build_loops(runner):
    val_loop = ValLoop(runner)
    if runner.loop_type == 'IterBasedTrainLoop':
        train_loop = IterBasedTrainLoop(runner, val_loop)
    else:
        train_loop = EpochBasedTrainLoop(runner, val_loop)

    runner.val_loop = val_loop
    runner.train_loop = train_loop
    runner._train_loop = train_loop
    return train_loop, val_loop


def maybe_resume_or_load(runner,
                         resume_from,
                         load_from,
                         resume_strict: bool,
                         load_strict: bool):
    if resume_from is not None:
        runner.logger.info(f'Resuming training from: {resume_from}')
        checkpoint = load_model_checkpoint(
            runner.model,
            resume_from,
            map_location='cpu',
            strict=resume_strict,
            logger=runner.logger)

        if checkpoint.get('optim_wrapper') is None:
            raise KeyError(
                f'Resume checkpoint {resume_from} is missing optim_wrapper state.')
        runner.optim_wrapper.load_state_dict(checkpoint['optim_wrapper'])

        if runner.param_schedulers and checkpoint.get('scheduler') is None:
            raise KeyError(
                f'Resume checkpoint {resume_from} is missing scheduler state.')
        load_scheduler_state_dict(runner.param_schedulers,
                                  checkpoint.get('scheduler'), resume_from)

        if checkpoint.get('message_hub') is None:
            raise KeyError(
                f'Resume checkpoint {resume_from} is missing message_hub state.')
        runner.message_hub.load_state_dict(checkpoint['message_hub'])

        if 'epoch' not in checkpoint or 'iter' not in checkpoint:
            raise KeyError(
                f'Resume checkpoint {resume_from} must contain explicit '
                'epoch and iter fields.')
        runner.start_epoch = int(checkpoint['epoch'])
        runner.global_iter = int(checkpoint['iter'])
        runner.logger.info(
            f'Resumed at epoch={runner.start_epoch}, iter={runner.global_iter}')
        runner._resume = True
        return checkpoint

    if load_from is not None:
        runner.logger.info(f'Loading pretrained weights from: {load_from}')
        return load_model_checkpoint(
            runner.model,
            load_from,
            map_location='cpu',
            strict=load_strict,
            logger=runner.logger)

    return None


def run_training(runner, train_loop, loaded_checkpoint) -> None:
    if loaded_checkpoint is not None:
        runner.hook_manager.call(
            'after_load_checkpoint', runner, checkpoint=loaded_checkpoint)

    runner.message_hub.update_info('epoch', runner.start_epoch)
    runner.message_hub.update_info('iter', runner.global_iter)
    runner.epoch = runner.start_epoch
    runner.iter = runner.global_iter
    runner.hook_manager.call('before_train', runner)

    runner.logger.info(f'Starting training loop with type: {runner.loop_type}')

    train_loop.run()

    runner.hook_manager.call('after_train', runner)
    runner.hook_manager.call('after_run', runner)
    runner.logger.info('Training completed!')


def main():
    args = parse_args()
    cfg = Config.fromfile(args.config)

    requested_distributed = args.launcher != 'none'
    setup_runtime_environment(cfg, requested_distributed)
    distributed = init_distributed_runtime(cfg, args.launcher)
    timestamp = sync_string_from_rank0(
        datetime.now().strftime('%Y%m%d_%H%M%S'), distributed)
    cfg.device = resolve_device(str(args.device), distributed)
    prepare_work_dir(cfg, args.config, args.work_dir, timestamp)
    barrier(distributed)

    logger = build_logger(cfg.work_dir)
    log_runtime_start(logger, args.config, cfg, distributed)
    setup_randomness(cfg, logger)

    resume_from, load_from, resume_strict, load_strict = (
        resolve_checkpoint_options(args, cfg))

    train_loop_cfg = resolve_train_loop_cfg(cfg)
    model = build_runtime_model(cfg, logger, resume_from, distributed)
    train_dataset, val_dataset, train_dataloader, val_dataloader = (
        build_data_loaders(cfg, train_loop_cfg['type'], logger))

    runner = build_runtime_context(
        cfg=cfg,
        logger=logger,
        model=model,
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        train_loop_cfg=train_loop_cfg,
        timestamp=timestamp)
    train_loop, _ = build_loops(runner)

    runner.hook_manager.call('before_run', runner)
    loaded_checkpoint = maybe_resume_or_load(
        runner=runner,
        resume_from=resume_from,
        load_from=load_from,
        resume_strict=resume_strict,
        load_strict=load_strict)
    run_training(runner, train_loop, loaded_checkpoint)


if __name__ == '__main__':
    try:
        main()
    finally:
        if (torch.distributed.is_available()
                and torch.distributed.is_initialized()):
            torch.distributed.destroy_process_group()