unitorch.cli.tasks.supervised

SupervisedTask

Tip

core/task/supervised is the section for configuration of SupervisedTask.

Standard supervised learning task with optional DDP, AMP, and EMA support.

Source code in src/unitorch/cli/tasks/supervised.py

def __init__(
    self,
    configure,
    model,
    datasets,
    local_rank: int = -1,  # GPU index for distributed training; -1 for single-GPU
    seed: int = 1123,  # global random seed for reproducibility
    cpu_offload: bool = False,  # keep model on CPU (e.g. for CPU-only environments)
):
    set_seed(seed)
    self.n_gpu = 1 if torch.cuda.is_available() else 0
    if dist.is_initialized():
        self.n_gpu = dist.get_world_size()

    self.config = configure
    self.model = model
    self.datasets = datasets
    self.local_rank = local_rank

    if self.local_rank != -1:
        torch.cuda.set_device(self.local_rank)

    if torch.cuda.is_available() and not cpu_offload:
        self.model = self.model.cuda()

    self.best_score = -np.inf

n_gpu instance-attribute

n_gpu = 1 if is_available() else 0

config instance-attribute

config = configure

model instance-attribute

model = model

datasets instance-attribute

datasets = datasets

local_rank instance-attribute

local_rank = local_rank

best_score instance-attribute

best_score = -inf

from_config classmethod

from_config(config, **kwargs)

Source code in src/unitorch/cli/tasks/supervised.py

@classmethod
@config_defaults_init("core/task/supervised")
def from_config(cls, config, **kwargs):
    try:
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
    except Exception:
        logging.info("PyTorch is not in distributed mode")

    config.set_default_section("core/task/supervised")

    model = config.getoption("model", None)
    dataset = config.getoption("dataset", None)

    if model is not None:
        model = init_registered_module(model, config, registered_model)
    if dataset is not None:
        dataset = init_registered_module(dataset, config, registered_dataset)

    return dict(
        configure=config,
        model=model,
        datasets=dataset,
        local_rank=config.getdefault("core/cli", "local_rank", get_local_rank()),
        cpu_offload=config.getoption("cpu_offload", False),
    )

train

train(
    optim: str,
    loss_fn: str,
    score_fn: str,
    monitor_fns: Optional[Union[str, List[str]]] = None,
    scheduler: Optional[str] = None,
    from_ckpt_dir: str = "./from_ckpt",
    to_ckpt_dir: str = "./to_ckpt",
    train_batch_size: int = 128,
    dev_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
    save_optimizer: bool = True,
    save_scheduler: bool = True,
    save_checkpoint: str = "default",
    log_freq: int = 100,
    ckpt_freq: int = 10000,
    grad_acc_step: int = 1,
    max_grad_norm: float = 1.0,
    num_training_samples: int = 1000000000,
    epochs: int = 5,
    use_ema: bool = False,
    ema_decay: float = 0.9999,
    ema_tau: int = 2000,
    use_amp: bool = True,
)

Source code in src/unitorch/cli/tasks/supervised.py

@config_defaults_method("core/task/supervised")
def train(
    self,
    optim: str,  # registered optimizer name
    loss_fn: str,  # registered loss function name
    score_fn: str,  # registered scoring function name
    monitor_fns: Optional[
        Union[str, List[str]]
    ] = None,  # extra metrics logged at checkpoints
    scheduler: Optional[str] = None,  # registered LR scheduler name
    from_ckpt_dir: str = "./from_ckpt",  # directory to load pretrained weights from
    to_ckpt_dir: str = "./to_ckpt",  # directory to write checkpoints to
    train_batch_size: int = 128,  # per-GPU batch size for training
    dev_batch_size: int = 128,  # per-GPU batch size for validation
    pin_memory: bool = True,  # pin DataLoader memory for faster GPU transfer
    num_workers: int = 4,  # DataLoader worker processes
    save_optimizer: bool = True,  # include optimizer state in checkpoints
    save_scheduler: bool = True,  # include scheduler state in checkpoints
    save_checkpoint: str = "default",  # checkpoint policy: default/best/latest/every/all
    log_freq: int = 100,  # log training loss every N steps
    ckpt_freq: int = 10000,  # save checkpoint every N steps
    grad_acc_step: int = 1,  # gradient accumulation steps before optimizer update
    max_grad_norm: float = 1.0,  # gradient clipping max norm
    num_training_samples: int = 1_000_000_000,  # fallback total samples for iterable datasets
    epochs: int = 5,  # total training epochs
    use_ema: bool = False,  # maintain an EMA shadow model for evaluation
    ema_decay: float = 0.9999,  # EMA decay factor
    ema_tau: int = 2000,  # EMA warm-up steps
    use_amp: bool = True,  # enable automatic mixed precision (FP16)
):
    if self.local_rank in [-1, 0]:
        os.makedirs(to_ckpt_dir, exist_ok=True)

    if loss_fn is not None:
        loss_fn = init_registered_module(loss_fn, self.config, registered_loss)
    if score_fn is not None:
        score_fn = init_registered_module(score_fn, self.config, registered_score)
    if monitor_fns is not None:
        monitor_fns = [
            init_registered_module(fn, self.config, registered_score)
            for fn in monitor_fns
            if fn in registered_score
        ]

    if optim is not None and self.model is not None:
        optim = init_registered_module(
            optim,
            self.config,
            registered_optim,
            params=filter(lambda p: p.requires_grad, self.model.parameters()),
        )

    # Load pretrained weights, then resume from latest checkpoint if available
    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)
        optim.from_checkpoint(from_ckpt_dir, weight_name="pytorch_optim.bin")
    if os.path.exists(to_ckpt_dir):
        self.model.from_checkpoint(
            to_ckpt_dir, weight_name="pytorch_model_latest.bin"
        )
        optim.from_checkpoint(to_ckpt_dir, weight_name="pytorch_optim_latest.bin")

    info_path = os.path.join(to_ckpt_dir, "info.json")
    if os.path.exists(info_path):
        with open(info_path) as f:
            info = json.load(f)
    else:
        info = {}

    global_epoch = info.get("global_epoch", 0)
    global_step = info.get("global_step", 0)
    self.best_score = info.get("best_score", self.best_score)
    logging.info("best score so far: %s", self.best_score)

    self.ema_model = None
    if use_ema:
        self.ema_model = ExponentialMovingAverage(
            self.model,
            decay=ema_decay,
            tau=ema_tau,
            num_steps=info.get("num_ema_steps", 0),
        )
        if os.path.exists(from_ckpt_dir):
            self.ema_model.from_checkpoint(
                from_ckpt_dir, weight_name="pytorch_ema_model.bin"
            )
        if os.path.exists(to_ckpt_dir):
            self.ema_model.from_checkpoint(
                to_ckpt_dir, weight_name="pytorch_ema_model_latest.bin"
            )

    for name, param in self.model.named_parameters():
        logging.debug(
            "%s: trainable=%s dtype=%s shape=%s device=%s",
            name,
            param.requires_grad,
            param.dtype,
            param.shape,
            param.device,
        )

    global_rank = -1
    if self.n_gpu > 1:
        self.model = nn.parallel.DistributedDataParallel(
            self.model,
            device_ids=[self.local_rank],
            output_device=self.local_rank,
            find_unused_parameters=False,
            broadcast_buffers=False,
        )
        global_rank = dist.get_rank()

    train_sampler = DistributedSkipSampler if self.n_gpu > 1 else RandomSkipSampler
    dev_sampler = DistributedSampler if self.n_gpu > 1 else SequentialSampler

    dataset_train = self.datasets.get("train")
    dataset_dev = self.datasets.get("dev")

    iter_train = DataLoader(
        dataset_train,
        sampler=(
            train_sampler(dataset_train)
            if not isinstance(dataset_train, Iterable)
            else None
        ),
        batch_size=train_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )
    iter_dev = DataLoader(
        dataset_dev,
        sampler=(
            dev_sampler(dataset_dev)
            if not isinstance(dataset_dev, Iterable)
            else None
        ),
        batch_size=dev_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    if scheduler is not None:
        if not isinstance(dataset_train, Iterable):
            num_training_steps = int(
                epochs
                * len(dataset_train)
                // train_batch_size
                // max(1, self.n_gpu)
                // grad_acc_step
            )
        else:
            num_training_steps = int(
                epochs
                * num_training_samples
                // train_batch_size
                // max(1, self.n_gpu)
                // grad_acc_step
            )
        scheduler = init_registered_module(
            scheduler,
            self.config,
            registered_scheduler,
            optimizer=optim,
            num_training_steps=num_training_steps,
        )

    if scheduler is not None and os.path.exists(to_ckpt_dir):
        scheduler.from_checkpoint(
            to_ckpt_dir, weight_name="pytorch_scheduler_latest.bin"
        )

    # AMP gradient scaler; only created when use_amp=True
    scaler = torch.amp.GradScaler("cuda") if use_amp else None

    def _optimizer_step():
        """Unscale gradients (if AMP), clip, then step the optimizer."""
        if scaler is not None:
            scaler.unscale_(optim)
        nn.utils.clip_grad_norm_(self.model.parameters(), max_grad_norm)
        if scaler is not None:
            scaler.step(optim)
            scaler.update()
        else:
            optim.step()
        if scheduler is not None:
            scheduler.step()
        optim.zero_grad()
        if use_ema and self.ema_model is not None:
            base = self.model.module if self.n_gpu > 1 else self.model
            self.ema_model.step(base)

    def _snapshot(epoch, step):
        return save_snapshot(
            self.model.module if self.n_gpu > 1 else self.model,
            to_ckpt_dir,
            iter_dev,
            score_fn,
            monitor_fns,
            optim=optim if save_optimizer else None,
            scheduler=scheduler if save_scheduler else None,
            save_checkpoint=save_checkpoint,
            ema_model=self.ema_model if use_ema else None,
            best_score=self.best_score,
            info_path=info_path,
            local_rank=self.local_rank,
            global_epoch=epoch,
            global_step=step,
        )

    log_loss = 0.0
    dev_epoch = 0

    for e in range(epochs):
        torch.cuda.empty_cache()
        if e < global_epoch:
            continue

        if hasattr(dataset_train, "set_epoch"):
            dataset_train.set_epoch(e)
        if hasattr(dataset_train, "set_skip_step"):
            dataset_train.set_skip_step(global_step * train_batch_size)
        if hasattr(iter_train.sampler, "set_epoch"):
            iter_train.sampler.set_epoch(e)
        if hasattr(iter_train.sampler, "set_skip_step"):
            iter_train.sampler.set_skip_step(global_step * train_batch_size)

        self.model.train()
        is_update_step = True

        for step, (inputs, targets) in enumerate(iter_train):
            step = step + global_step
            is_update_step = False

            if torch.cuda.is_available():
                inputs = inputs.cuda()
                targets = targets.cuda()

            with torch.autocast(
                device_type="cuda" if torch.cuda.is_available() else "cpu",
                enabled=use_amp,
            ):
                outputs = self.model(**inputs.dict())
                loss = (
                    outputs.loss
                    if isinstance(outputs, LossOutputs)
                    else loss_fn(outputs=outputs, targets=targets)
                ) / grad_acc_step

            if scaler is not None:
                scaler.scale(loss).backward()
            else:
                loss.backward()

            log_loss += loss.item() * grad_acc_step

            if (step + 1) % grad_acc_step == 0:
                is_update_step = True
                _optimizer_step()

            if (step + 1) % log_freq == 0 and global_rank in [-1, 0]:
                avg_loss = log_loss / log_freq
                logging.info("epoch %d step %d: train/loss=%.6f", e, step, avg_loss)
                if wandb.is_available():
                    wandb.log({"epoch": e, "step": step, "train/loss": avg_loss})
                log_loss = 0.0

            if (step + 1) % ckpt_freq == 0:
                if hasattr(dataset_dev, "set_epoch"):
                    dataset_dev.set_epoch(dev_epoch)
                if hasattr(iter_dev.sampler, "set_epoch"):
                    iter_dev.sampler.set_epoch(dev_epoch)
                dev_epoch += 1
                self.best_score = _snapshot(e, step + 1)

        # Flush any remaining accumulated gradients at epoch end
        if not is_update_step:
            _optimizer_step()

        log_loss = 0.0

        if hasattr(dataset_dev, "set_epoch"):
            dataset_dev.set_epoch(dev_epoch)
        if hasattr(iter_dev.sampler, "set_epoch"):
            iter_dev.sampler.set_epoch(dev_epoch)
        dev_epoch += 1

        global_step = 0
        self.best_score = _snapshot(e + 1, 0)

eval

eval(
    monitor_fns: Union[str, List[str]],
    from_ckpt_dir: str = "./from_ckpt",
    dev_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
)

Source code in src/unitorch/cli/tasks/supervised.py

@torch.no_grad()
@config_defaults_method("core/task/supervised")
def eval(
    self,
    monitor_fns: Union[str, List[str]],  # list of registered scoring function names
    from_ckpt_dir: str = "./from_ckpt",  # directory to load model weights from
    dev_batch_size: int = 128,  # per-GPU batch size for evaluation
    pin_memory: bool = True,
    num_workers: int = 4,
):
    monitor_fns = [
        init_registered_module(fn, self.config, registered_score)
        for fn in monitor_fns
        if fn in registered_score
    ]

    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)

    global_rank = -1
    if self.n_gpu > 1:
        self.model = nn.parallel.DistributedDataParallel(
            self.model,
            device_ids=[self.local_rank],
            output_device=self.local_rank,
            find_unused_parameters=False,
            broadcast_buffers=False,
        )
        global_rank = dist.get_rank()

    dev_sampler = DistributedSampler if self.n_gpu > 1 else SequentialSampler
    dataset_dev = self.datasets.get("dev")
    iter_dev = DataLoader(
        dataset_dev,
        sampler=(
            dev_sampler(dataset_dev)
            if not isinstance(dataset_dev, Iterable)
            else None
        ),
        batch_size=dev_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    results = infer(self.model.module if self.n_gpu > 1 else self.model, iter_dev)
    if global_rank in [-1, 0]:
        monitor(
            outputs=results.outputs,
            targets=results.targets,
            monitor_fns=monitor_fns,
        )

infer

infer(
    postprocess_fn: str,
    writer: str,
    test_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
    max_size: int = 10000,
    from_ckpt_dir: str = "./from_ckpt",
    output_header: Optional[List] = None,
    output_path: str = "./output.txt",
    postprocess_workers: int = 2,
)

Source code in src/unitorch/cli/tasks/supervised.py

@torch.no_grad()
@config_defaults_method("core/task/supervised")
def infer(
    self,
    postprocess_fn: str,  # registered postprocessing function name
    writer: str,  # registered writer name for output serialisation
    test_batch_size: int = 128,  # per-GPU batch size for inference
    pin_memory: bool = True,
    num_workers: int = 4,
    max_size: int = 10000,  # maximum queue depth for async postprocessing
    from_ckpt_dir: str = "./from_ckpt",  # directory to load model weights from
    output_header: Optional[
        List
    ] = None,  # column names to copy from raw dataset into output
    output_path: str = "./output.txt",  # file path for inference results
    postprocess_workers: int = 2,  # number of parallel postprocessing workers
):
    assert self.n_gpu <= 1, "inference only supports single-GPU mode"
    assert writer is not None

    output_dir = os.path.dirname(output_path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    if postprocess_fn is not None:
        postprocess_fn = init_registered_process(postprocess_fn, self.config)

    writer = init_registered_module(
        writer, self.config, registered_writer, output_file=output_path
    )
    skip_step = writer.skip_n_samples

    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)

    sampler = SequentialSkipSampler if skip_step > 0 else SequentialSampler
    dataset_test = self.datasets.get("test")

    iter_test = DataLoader(
        dataset_test,
        sampler=(
            sampler(dataset_test)
            if not isinstance(dataset_test, Iterable)
            else None
        ),
        batch_size=test_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    if skip_step > 0:
        if hasattr(dataset_test, "set_skip_step"):
            dataset_test.set_skip_step(skip_step)
        if hasattr(iter_test.sampler, "set_skip_step"):
            iter_test.sampler.set_skip_step(skip_step)

    # Build a parallel loader for raw dataset metadata (images, text) when available
    iter_data = None
    if hasattr(dataset_test, "dataset"):
        data_info = DatasetFeature(dataset_test.dataset)
        iter_data = DataLoader(
            deepcopy(data_info),
            sampler=(
                sampler(data_info)
                if not isinstance(dataset_test, Iterable)
                else None
            ),
            batch_size=test_batch_size,
            shuffle=False,
            pin_memory=pin_memory,
            num_workers=num_workers,
            collate_fn=None,
        )
        if skip_step > 0 and hasattr(iter_data.sampler, "set_skip_step"):
            iter_data.sampler.set_skip_step(skip_step)

    self.model.eval()
    start = time.time()

    data_queue = Queue(maxsize=max_size)
    msg_queue = Queue(maxsize=max_size)
    postprocess_list = [
        PostProcess(postprocess_fn, data_queue, msg_queue)
        for _ in range(postprocess_workers)
    ]
    for p in postprocess_list:
        p.start()

    io_process = IOProcess(msg_queue, writer=writer)
    io_process.start()

    if iter_data is None:
        for step, (inputs, _) in enumerate(iter_test):
            if torch.cuda.is_available():
                inputs = inputs.cuda()
            outputs = self.model(**inputs.dict()).cpu()
            data_queue.put((step, outputs))
    else:
        for step, ((inputs, _), raw_info) in enumerate(zip(iter_test, iter_data)):
            if torch.cuda.is_available():
                inputs = inputs.cuda()
            outputs = self.model(**inputs.dict()).cpu()
            if output_header is not None:
                raw_info = {k: raw_info[k] for k in output_header if k in raw_info}
                outputs.from_pandas(pd.DataFrame(raw_info))
            data_queue.put((step, outputs))

    data_queue.put((-1, GENERATE_FINISHED))
    for p in postprocess_list:
        p.join()

    msg_queue.put((-1, GENERATE_FINISHED))
    io_process.join()

    elapsed_ms = (time.time() - start) * 1000
    throughput = (len(dataset_test) - skip_step) / elapsed_ms * 1000
    logging.info("%.2f ms | %.2f samples/s", elapsed_ms, throughput)