tokenizer.py

import os
import logging
import math
import time
import json
import datetime
import numpy as np

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from timm import create_model

import glob as glob_module

import utils
from model.shape_recon import shape_recon
from sentence.shapeefm_reshaping import ShapeEFMReshaping

logger = logging.getLogger('__main__')


def setup_logging(rank):
    root = logging.getLogger()
    root.handlers.clear()
    fmt = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    if rank == 0:
        sh = logging.StreamHandler()
        sh.setFormatter(fmt)
        root.addHandler(sh)
        fh = logging.FileHandler('tokenizer.log')
        fh.setFormatter(fmt)
        root.addHandler(fh)
        root.setLevel(logging.INFO)
    else:
        root.setLevel(logging.WARNING)


def has_checkpoint(output_dir):
    """Check if any checkpoint exists in the directory."""
    return bool(glob_module.glob(os.path.join(output_dir, "checkpoint*.pth")))


def create_shape_recon_model(args):
    return create_model(
        args.tokenizer_model,
        pretrained=False,
        code_num=args.codebook_size,
    )


def train_one_epoch(model, data_loader, optimizer, device, epoch,
                    loss_scaler, clip_grad=0, log_writer=None,
                    start_steps=None, lr_schedule_values=None, args=None):
    model.train()
    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))
    metric_logger.add_meter("min_lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))
    header = f"Epoch: [{epoch}]"

    _model = model.module if hasattr(model, "module") else model
    if hasattr(_model, "quantize"):
        try:
            _model.quantize.reset_cluster_size(device)
        except (AttributeError, RuntimeError):
            pass

    for step, batch in enumerate(metric_logger.log_every(data_loader, 100, header)):
        it = min(start_steps + step, len(lr_schedule_values) - 1) if lr_schedule_values is not None else start_steps + step
        if lr_schedule_values is not None:
            for param_group in optimizer.param_groups:
                param_group["lr"] = lr_schedule_values[it] * param_group.get("lr_scale", 1.0)

        ECG = batch.float().to(device, non_blocking=True)
        with torch.amp.autocast(device_type=device.type, enabled=True):
            loss, log_loss = model(ECG)

        loss_value = loss.item()
        if not math.isfinite(loss_value):
            raise FloatingPointError(f"Loss is {loss_value}, stopping training")

        optimizer.zero_grad()
        is_second_order = hasattr(optimizer, "is_second_order") and optimizer.is_second_order
        grad_norm = loss_scaler(loss, optimizer, clip_grad=clip_grad,
                                parameters=model.parameters(), create_graph=is_second_order)
        loss_scale_value = loss_scaler.state_dict()["scale"]
        torch.cuda.synchronize()

        metric_logger.update(loss=loss_value)
        new_log_loss = {k.split("/")[-1]: v for k, v in log_loss.items() if k != "total_loss"}
        metric_logger.update(**new_log_loss)

        min_lr = min(g["lr"] for g in optimizer.param_groups)
        max_lr = max(g["lr"] for g in optimizer.param_groups)
        metric_logger.update(lr=max_lr, min_lr=min_lr)
        weight_decay_value = next((g["weight_decay"] for g in optimizer.param_groups if g["weight_decay"] > 0), None)
        metric_logger.update(weight_decay=weight_decay_value, grad_norm=grad_norm)

        if log_writer is not None:
            log_writer.update(**new_log_loss, head="train/loss")
            log_writer.update(lr=max_lr, min_lr=min_lr, weight_decay=weight_decay_value,
                              grad_norm=grad_norm, loss_scale=loss_scale_value, head="opt")
            log_writer.set_step()

    metric_logger.synchronize_between_processes()
    if dist.get_rank() == 0:
        logger.info("Averaged stats: %s", metric_logger)

    stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()}
    if hasattr(_model, "quantize"):
        try:
            cluster_size = _model.quantize._codebook.cluster_size
        except AttributeError:
            cluster_size = _model.quantize.cluster_size
        zero_cnt = (cluster_size == 0).sum().item()
        stats["unused_code"] = zero_cnt
        if dist.get_rank() == 0:
            logger.info(f"Unused code in codebook: {zero_cnt}")
    return stats


@torch.no_grad()
def evaluate(data_loader, model, device, log_writer=None, epoch=None, args=None):
    model.eval()
    metric_logger = utils.MetricLogger(delimiter="  ")
    header = "Validation:"

    _model = model.module if hasattr(model, "module") else model
    if hasattr(_model, "quantize"):
        try:
            _model.quantize.reset_cluster_size(device)
        except (AttributeError, RuntimeError):
            pass

    codebook_vis = getattr(args, 'codebook_visualization', False)
    if codebook_vis:
        data_list, quantize_list, labels_list = [], [], []

    for step, batch in enumerate(metric_logger.log_every(data_loader, 10, header)):
        ECG = batch.float().to(device, non_blocking=True)
        loss, log_loss = model(ECG)
        metric_logger.update(loss=loss.item())
        new_log_loss = {k.split("/")[-1]: v for k, v in log_loss.items() if k != "total_loss"}
        metric_logger.update(**new_log_loss)

        if codebook_vis:
            tokens = _model.get_tokens(ECG)
            data_list.extend(ECG.cpu().detach().numpy().reshape(-1, ECG.size(-1)))
            quantize_list.extend(tokens["quantize"].cpu().detach().numpy().reshape(-1, tokens["quantize"].size(-1)))
            labels_list.extend(tokens["token"].cpu().detach().numpy().reshape(-1))

    if codebook_vis and dist.get_rank() == 0:
        flat_data, flat_quantize, flat_labels = np.array(data_list), np.array(quantize_list), np.array(labels_list)
        logger.info(f"Codebook vis shapes: data={flat_data.shape}, quantize={flat_quantize.shape}, labels={flat_labels.shape}")

    metric_logger.synchronize_between_processes()
    if dist.get_rank() == 0:
        logger.info("Averaged stats: %s", metric_logger)

    stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()}
    if hasattr(_model, "quantize"):
        try:
            cluster_size = _model.quantize._codebook.cluster_size
        except AttributeError:
            cluster_size = _model.quantize.cluster_size
        zero_cnt = (cluster_size == 0).sum().item()
        stats["unused_code"] = zero_cnt
        if dist.get_rank() == 0:
            logger.info(f"Unused code in codebook: {zero_cnt}")
    return stats


def count_parameters(model, model_name):
    n_learnable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    n_fixed = sum(p.numel() for p in model.parameters() if not p.requires_grad)
    if dist.get_rank() == 0:
        logger.info(f"{model_name}: {n_learnable / 1e6:.1f}M learnable, {n_fixed / 1e6:.1f}M fixed params")
    return n_learnable, n_fixed


def main():
    args = utils.get_args_shape_recon()
    utils.setup_distributed(args)
    rank = dist.get_rank()
    setup_logging(rank)
    opts = utils.setup_args(args)
    is_main = rank == 0

    base_output_dir = opts.output_dir
    base_lr = opts.lr

    # Create model once (shared across all datasets)
    model = create_shape_recon_model(opts).to(opts.device)
    if opts.tokenizer_weight:
        ckpt = torch.load(opts.tokenizer_weight, map_location="cpu", weights_only=False)
        model.load_state_dict(ckpt.get("model", ckpt), strict=False)
        if is_main:
            logger.info(f"Loaded initial weights from {opts.tokenizer_weight}")

    for part_name in ["encoder", "decoder"]:
        count_parameters(getattr(model, part_name), part_name)
    n_learnable, _ = count_parameters(model, "shape_recon")

    model = DDP(model, device_ids=[opts.gpu], find_unused_parameters=False)

    if is_main:
        logger.info(f"Sequential tokenizer training on: {opts.datasets}")
    start_time = time.time()

    for ds_idx, dataset_name in enumerate(opts.datasets):
        if is_main:
            logger.info(f"{'=' * 60}")
            logger.info(f"[{ds_idx + 1}/{len(opts.datasets)}] Training tokenizer on {dataset_name}")
            logger.info(f"{'=' * 60}")

        # Configure paths for this dataset
        utils.configure_dataset(opts, dataset_name)
        opts.output_dir = os.path.join(base_output_dir, dataset_name)
        os.makedirs(opts.output_dir, exist_ok=True)

        # Build data pipeline
        precompute = getattr(opts, 'precompute_data', True)
        train_loader, val_loader, train_num, val_num = \
            ShapeEFMReshaping(opts, precompute=precompute).run()
        if is_main:
            logger.info(f"Train: {train_num}, Val: {val_num}")

        num_steps_per_epoch = train_num // opts.batch_size // opts.world_size

        # Fresh optimizer & scheduler for each dataset
        optimizer = utils.create_optimizer(opts, model.module)
        loss_scaler = utils.NativeScaler()

        total_batch = opts.batch_size * opts.world_size
        scaled_lr = total_batch / 128 * base_lr
        if is_main:
            logger.info(f"LR = {scaled_lr:.8f}, Batch = {total_batch}, Steps/epoch = {num_steps_per_epoch}")

        lr_schedule_values = utils.cosine_scheduler(
            scaled_lr, opts.min_lr, opts.epochs,
            num_steps_per_epoch,
            warmup_epochs=opts.warmup_epochs,
            warmup_steps=opts.warmup_steps,
        )

        # Resume from this dataset's checkpoint, or load from previous dataset
        opts.start_epoch = 1
        opts.resume = ""
        if has_checkpoint(opts.output_dir):
            utils.auto_load_model(args=opts, model=model, model_without_ddp=model.module, optimizer=optimizer, loss_scaler=loss_scaler)
        elif ds_idx > 0:
            prev_ckpt = os.path.join(base_output_dir, opts.datasets[ds_idx - 1], "checkpoint.pth")
            if os.path.exists(prev_ckpt):
                ckpt = torch.load(prev_ckpt, map_location="cpu", weights_only=False)
                model.module.load_state_dict(ckpt["model"], strict=False)
                if is_main:
                    logger.info(f"Loaded weights from previous dataset: {prev_ckpt}")

        if opts.eval:
            evaluate(val_loader, model, opts.device, args=opts)
            continue

        # TensorBoard per dataset
        tb_dir = os.path.join(opts.log_dir, f'tb_{dataset_name}')
        os.makedirs(tb_dir, exist_ok=True)
        log_writer = utils.TensorboardLogger(tb_dir) if is_main else None

        for epoch in range(opts.start_epoch, opts.epochs + 1):
            train_loader.sampler.set_epoch(epoch - 1)
            if log_writer is not None:
                log_writer.set_step((epoch - 1) * num_steps_per_epoch)

            train_stats = train_one_epoch(
                model, train_loader, optimizer, opts.device, epoch,
                loss_scaler, opts.clip_grad, log_writer=log_writer,
                start_steps=(epoch - 1) * num_steps_per_epoch,
                lr_schedule_values=lr_schedule_values, args=opts,
            )

            if opts.output_dir and epoch % opts.save_ckpt_freq == 0:
                utils.save_model(
                    args=opts, epoch=epoch, model=model,
                    model_without_ddp=model.module,
                    optimizer=optimizer, loss_scaler=loss_scaler,
                    save_ckpt_freq=opts.save_ckpt_freq,
                )

            if val_loader is not None and epoch % opts.val_freq == 0:
                test_stats = evaluate(val_loader, model, opts.device, log_writer, epoch, args=opts)
                if is_main:
                    logger.info(f"Val loss: {test_stats['loss']:.4f}")
                if log_writer is not None:
                    log_writer.update(**test_stats, head="val/loss")

                log_stats = {
                    **{f"train_{k}": v for k, v in train_stats.items()},
                    **{f"val_{k}": v for k, v in test_stats.items()},
                    "epoch": epoch, "dataset": dataset_name, "n_parameters": n_learnable,
                }
            else:
                log_stats = {
                    **{f"train_{k}": v for k, v in train_stats.items()},
                    "epoch": epoch, "dataset": dataset_name, "n_parameters": n_learnable,
                }

            if opts.output_dir and epoch % opts.val_freq == 0:
                if log_writer is not None:
                    log_writer.flush()
                with open(os.path.join(opts.output_dir, "log.txt"), "a") as f:
                    f.write(json.dumps(log_stats) + "\n")

        if is_main:
            logger.info(f"Completed tokenizer training on {dataset_name}")

    total_time = str(datetime.timedelta(seconds=int(time.time() - start_time)))
    if is_main:
        logger.info(f"Total training time: {total_time}")


if __name__ == '__main__':
    main()