genhpf 1.0.11__py3-none-any.whl

genhpf/trainer.py

@@ -0,0 +1,370 @@
+import logging
+import time
+from itertools import chain
+from typing import Any, Dict, List, Tuple
+
+import torch
+import torch.distributed as dist
+import torch.utils
+import torch.utils.data
+from omegaconf import OmegaConf
+from torch.optim import Adam
+
+from genhpf.configs import Config
+from genhpf.datasets import BaseDataset
+from genhpf.loggings import metrics
+from genhpf.utils import checkpoint_utils, distributed_utils, utils
+from genhpf.utils.file_io import PathManager
+
+logger = logging.getLogger(__name__)
+
+
+class Trainer(object):
+    def __init__(self, cfg: Config, model, criterion):
+        self.cfg = cfg
+        self.cuda = torch.cuda.is_available()
+        if self.cuda:
+            self.device = torch.device("cuda")
+        else:
+            self.device = torch.device("cpu")
+
+        self._criterion = criterion
+        self._model = model
+
+        self._criterion = self._criterion.to(device=self.device)
+        self._model = self._model.to(device=self.device)
+
+        self._num_updates = 0
+
+        self._optimizer = None
+        self._wrapped_criterion = None
+        self._wrapped_model = None
+
+        if self.cuda:
+            self.cuda_env = utils.CudaEnvironment()
+            if self.data_parallel_world_size > 1:
+                self.cuda_env_arr = distributed_utils.all_gather_list(
+                    self.cuda_env, group=distributed_utils.get_global_group()
+                )
+            else:
+                self.cuda_env_arr = [self.cuda_env]
+            if self.data_parallel_rank == 0:
+                utils.CudaEnvironment.pretty_print_cuda_env_list(self.cuda_env_arr)
+        else:
+            self.cuda_env = None
+            self.cuda_env_arr = None
+
+        metrics.log_start_time("wall", priority=790, round=0)
+
+        self._start_time = time.time()
+        self._previous_training_time = 0
+        self._cumulative_training_time = None
+
+    @property
+    def data_parallel_world_size(self):
+        if self.cfg.distributed_training.distributed_world_size == 1:
+            return 1
+        return distributed_utils.get_data_parallel_world_size()
+
+    @property
+    def data_parallel_process_group(self):
+        return distributed_utils.get_data_parallel_group()
+
+    @property
+    def data_parallel_rank(self):
+        if self.cfg.distributed_training.distributed_world_size == 1:
+            return 0
+        return distributed_utils.get_data_parallel_rank()
+
+    @property
+    def is_data_parallel_master(self):
+        return self.data_parallel_rank == 0
+
+    @property
+    def use_distributed_wrapper(self) -> bool:
+        return self.data_parallel_world_size > 1
+
+    @property
+    def criterion(self):
+        if self._wrapped_criterion is None:
+            if utils.has_parameters(self._criterion) and self.use_distributed_wrapper:
+                self._wrapped_criterion = distributed_utils.DistributedModel(
+                    self.cfg.distributed_training,
+                    self._criterion,
+                    process_group=self.data_parallel_process_group,
+                    device=self.device,
+                )
+            else:
+                self._wrapped_criterion = self._criterion
+        return self._wrapped_criterion
+
+    @property
+    def model(self):
+        if self._wrapped_model is None:
+            if self.use_distributed_wrapper:
+                self._wrapped_model = distributed_utils.DistributedModel(
+                    self.cfg.distributed_training,
+                    self._model,
+                    process_group=self.data_parallel_process_group,
+                    device=self.device,
+                )
+            else:
+                self._wrapped_model = self._model
+        return self._wrapped_model
+
+    @property
+    def optimizer(self):
+        if self._optimizer is None:
+            self._build_optimizer()
+        return self._optimizer
+
+    def _build_optimizer(self):
+        params = list(
+            filter(lambda p: p.requires_grad, chain(self.model.parameters(), self.criterion.parameters()))
+        )
+        self._optimizer = Adam(
+            params,
+            lr=self.cfg.optimization.lr,
+            betas=self.cfg.optimization.adam_betas,
+            eps=self.cfg.optimization.adam_eps,
+            weight_decay=self.cfg.optimization.weight_decay,
+        )
+
+    def state_dict(self):
+        pretraining_parameter_names = []
+        if hasattr(self.model, "get_pretraining_parameter_names"):
+            pretraining_parameter_names = self.model.get_pretraining_parameter_names()
+        model_state_dict = {
+            k: v for k, v in self.model.state_dict().items() if k not in pretraining_parameter_names
+        }
+
+        state_dict = {
+            "cfg": (
+                OmegaConf.to_container(self.cfg, resolve=True, enum_to_str=True)
+                if OmegaConf.is_config(self.cfg)
+                else self.cfg
+            ),
+            "model": model_state_dict,
+        }
+
+        return state_dict
+
+    def save_checkpoint(self, filename):
+        """Save all training state in a checkpoint file."""
+        logger.info(f"Saving checkpoint to {filename}")
+
+        state_dict = utils.move_to_cpu(self.state_dict())
+        if self.is_data_parallel_master:
+            checkpoint_utils.torch_persistent_save(state_dict, filename, async_write=False)
+        logger.info(f"Finished saving checkpoint to {filename}")
+
+    def load_checkpoint(self, filename) -> None:
+        """
+        Load all training state from a checkpoint file.
+        rank = 0 will load the checkpoint and broadcast it to other ranks.
+        """
+        logger.info(f"Loading checkpoint from {filename}")
+        is_distributed = self.data_parallel_world_size > 1
+        bexists = PathManager.isfile(filename)
+        if bexists:
+            if self.data_parallel_rank == 0:
+                state = checkpoint_utils.load_checkpoint_to_cpu(filename)
+            else:
+                state = None
+
+            if is_distributed:
+                state = distributed_utils.broadcast_object(
+                    state, src_rank=0, group=self.data_parallel_process_group, dist_device=self.device
+                )
+
+            # load model parameters
+            try:
+                self.model.load_state_dict(state["model"], strict=True)
+                # save memory for later steps
+                del state["model"]
+            except Exception:
+                raise Exception(
+                    f"Cannot load model parameters from checkpoint {filename}; "
+                    "please ensure that the architectures match."
+                )
+        else:
+            logger.info(f"No existing checkpoint found {filename}")
+
+    def get_train_iterator(
+        self,
+        dataset: BaseDataset,
+    ) -> Tuple[torch.utils.data.DataLoader, torch.utils.data.Sampler]:
+        """Return an DataLoader instance for training."""
+        batch_sampler = (
+            torch.utils.data.DistributedSampler(dataset, shuffle=True) if dist.is_initialized() else None
+        )
+        batch_iterator = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=self.cfg.dataset.batch_size,
+            shuffle=True if not dist.is_initialized() else False,
+            num_workers=self.cfg.dataset.num_workers if not self.cfg.common.debug else 0,
+            collate_fn=dataset.collator,
+            sampler=batch_sampler,
+        )
+        return batch_iterator, batch_sampler
+
+    def get_valid_iterator(
+        self,
+        dataset: BaseDataset,
+    ) -> Tuple[torch.utils.data.DataLoader, torch.utils.data.Sampler]:
+        """Return an DataLoader instance for validation."""
+        batch_sampler = (
+            torch.utils.data.DistributedSampler(dataset, shuffle=False) if dist.is_initialized() else None
+        )
+        batch_iterator = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=self.cfg.dataset.batch_size,
+            shuffle=False,
+            num_workers=self.cfg.dataset.num_workers if not self.cfg.common.debug else 0,
+            collate_fn=dataset.collator,
+            sampler=batch_sampler,
+        )
+        return batch_iterator, batch_sampler
+
+    @metrics.aggregate("train")
+    def train_step(self, sample):
+        """Do forward, backward and optimization steps."""
+
+        self._set_seed()
+        self.model.train()
+        self.criterion.train()
+        self.zero_grad()
+
+        metrics.log_start_time("train_wall", priority=800, round=0)
+
+        logging_outputs = []
+        sample = utils.prepare_sample(sample)
+
+        loss, sample_size, logging_output = self.criterion(self.model, sample)
+        if loss.item() > 0:
+            loss.backward()  # backward
+            self.optimizer.step()
+
+        logging_outputs.append(logging_output)
+
+        # emptying the CUDA cache after the first step can
+        # reduce the chance of OOM
+        if self.cuda and self.get_num_updates() == 0:
+            torch.cuda.empty_cache()
+
+        sample_size = float(sample_size)
+
+        if self._sync_stats():
+            train_time = self._local_cumulative_training_time()
+
+            logging_outputs, (sample_size, total_train_time) = self._aggregate_logging_outputs(
+                logging_outputs, sample_size, train_time
+            )
+
+            self._cumulative_training_time = total_train_time / self.data_parallel_world_size
+
+        logging_output = None
+        self.set_num_updates(self.get_num_updates() + 1)
+
+        if self.cuda and self.cuda_env is not None:
+            gb_used = torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024
+            torch.cuda.reset_peak_memory_stats()
+            gb_free = self.cuda_env.total_memory_in_GB - gb_used
+            metrics.log_scalar("gb_free", gb_free, priority=1500, round=1, weight=0)
+
+        # extract private logs (usually only for valid steps) before logging
+        logging_outputs = list(
+            map(lambda x: {key: x[key] for key in x if not key.startswith("_")}, logging_outputs)
+        )
+        # log stats
+        logging_output = self._reduce_and_log_stats(logging_outputs, sample_size)
+
+        metrics.log_stop_time("train_wall")
+        return logging_output
+
+    @metrics.aggregate("valid")
+    def valid_step(self, sample, subset=None):
+        """Do forward pass in evaluation mode."""
+        with torch.no_grad():
+            self.model.eval()
+            self.criterion.eval()
+
+            sample = utils.prepare_sample(sample)
+            _loss, sample_size, logging_output = self.criterion(self.model, sample)
+            logging_outputs = [logging_output]
+
+        # gather logging outputs from all replicas
+        if self.data_parallel_world_size > 1:
+            logging_outputs, (sample_size,) = self._aggregate_logging_outputs(
+                logging_outputs,
+                sample_size,
+                ignore=False,
+            )
+
+        # log validation stats
+        logging_output = self._reduce_and_log_stats(logging_outputs, sample_size)
+
+        return _loss, sample_size, logging_outputs
+
+    def zero_grad(self):
+        self.optimizer.zero_grad()
+
+    def get_num_updates(self):
+        return self._num_updates
+
+    def set_num_updates(self, num_updates):
+        self._num_updates = num_updates
+        metrics.log_scalar("num_updates", self._num_updates, weight=0, priority=200)
+
+    def cumulative_training_time(self):
+        if self._cumulative_training_time is None:
+            return self._local_cumulative_training_time()
+        else:
+            return self._cumulative_training_time
+
+    def _local_cumulative_training_time(self):
+        return time.time() - self._start_time + self._previous_training_time
+
+    def _set_seed(self):
+        seed = self.cfg.common.seed + self.get_num_updates()
+        utils.set_torch_seed(seed)
+
+    def _sync_stats(self):
+        if self.data_parallel_world_size == 1:
+            return False
+        else:
+            return True
+
+    def _aggregate_logging_outputs(
+        self, logging_outputs: List[Dict[str, Any]], *extra_stats_to_sum, ignore=False
+    ):
+        return self._all_gather_list_sync(logging_outputs, *extra_stats_to_sum, ignore=ignore)
+
+    def _all_gather_list_sync(self, logging_outputs: List[Dict[str, Any]], *extra_stats_to_sum, ignore=False):
+        if ignore:
+            logging_outputs = []
+        results = list(
+            zip(
+                *distributed_utils.all_gather_list(
+                    [logging_outputs] + list(extra_stats_to_sum),
+                    max_size=getattr(self.cfg.common, "all_gather_list_size", 32768),
+                    group=self.data_parallel_process_group,
+                )
+            )
+        )
+        logging_outputs, extra_stats_to_sum = results[0], results[1:]
+        logging_outputs = list(chain.from_iterable(logging_outputs))
+        extra_stats_to_sum = [sum(s) for s in extra_stats_to_sum]
+        return logging_outputs, extra_stats_to_sum
+
+    def _reduce_and_log_stats(self, logging_outputs, sample_size):
+        metrics.log_speed("ups", 1.0, priority=100, round=2)
+
+        with metrics.aggregate() as agg:
+            if logging_outputs is not None:
+                self.criterion.__class__.reduce_metrics(logging_outputs)
+                del logging_outputs
+
+            logging_output = agg.get_smoothed_values()
+            logging_output["sample_size"] = sample_size
+            return logging_output

genhpf/utils/checkpoint_utils.py

@@ -0,0 +1,171 @@
+import collections
+import logging
+import os
+import re
+import traceback
+
+import torch
+
+from genhpf.configs import CheckpointConfig
+from genhpf.utils.file_io import PathManager
+
+logger = logging.getLogger(__name__)
+
+
+def save_checkpoint(cfg: CheckpointConfig, trainer, epoch: int, val_loss: float):
+    from genhpf.loggings import meters
+
+    # only one worker should attempt to create the required dir
+    if trainer.data_parallel_rank == 0:
+        os.makedirs(cfg.save_dir, exist_ok=True)
+
+    prev_best = getattr(save_checkpoint, "best", val_loss)
+    if val_loss is not None:
+        best_function = max if cfg.maximize_best_checkpoint_metric else min
+        save_checkpoint.best = best_function(val_loss, prev_best)
+
+    if cfg.no_save:
+        return
+
+    if not trainer.is_data_parallel_master:
+        return
+
+    write_timer = meters.StopwatchMeter()
+    write_timer.start()
+
+    updates = trainer.get_num_updates()
+
+    logger.info(f"preparing to save checkpoint for epoch {epoch} @ {updates} updates")
+
+    def is_better(a, b):
+        return a >= b if cfg.maximize_best_checkpoint_metric else a <= b
+
+    checkpoint_conds = collections.OrderedDict()
+    checkpoint_conds["checkpoint_{}.pt".format(epoch)] = epoch % cfg.save_interval == 0
+    checkpoint_conds["checkpoint_best.pt"] = val_loss is not None and (
+        not hasattr(save_checkpoint, "best") or is_better(val_loss, save_checkpoint.best)
+    )
+
+    checkpoint_conds["checkpoint_last.pt"] = not cfg.no_last_checkpoints
+
+    checkpoints = [os.path.join(cfg.save_dir, fn) for fn, cond in checkpoint_conds.items() if cond]
+    if len(checkpoints) > 0:
+        trainer.save_checkpoint(checkpoints[0])
+        for cp in checkpoints[1:]:
+            assert PathManager.copy(
+                checkpoints[0], cp, overwrite=True
+            ), f"Failed to copy {checkpoints[0]} to {cp}"
+
+        write_timer.stop()
+        logger.info(
+            "Save checkpoint {} (epoch {} @ {} updates, score {}) (writing took {} seconds)".format(
+                checkpoints[0], epoch, updates, val_loss, write_timer.sum
+            )
+        )
+
+    if cfg.keep_last_epochs > 0:
+        # remove old epoch checkpoints; checkpoints are sorted in descending order
+        checkpoints = checkpoint_paths(cfg.save_dir, pattern=r"checkpoint_(\d+)\.pt")
+        for old_chk in checkpoints[cfg.keep_last_epochs :]:
+            if os.path.lexists(old_chk):
+                os.remove(old_chk)
+
+
+def torch_persistent_save(obj, filename, async_write: bool = False):
+    if async_write:
+        with PathManager.opena(filename, "wb") as f:
+            _torch_persistent_save(obj, f)
+    else:
+        if PathManager.supports_rename(filename):
+            # do atmoic save
+            with PathManager.open(filename + ".tmp", "wb") as f:
+                _torch_persistent_save(obj, f)
+            PathManager.rename(filename + ".tmp", filename)
+
+
+def _torch_persistent_save(obj, f):
+    if isinstance(f, str):
+        with PathManager.open(f, "wb") as h:
+            _torch_persistent_save(obj, h)
+        return
+    for i in range(3):
+        try:
+            return torch.save(obj, f)
+        except Exception:
+            if i == 2:
+                logger.error(traceback.format_exc())
+
+
+def checkpoint_paths(path, pattern=r"checkpoint(\d+)\.pt", keep_match=False):
+    """Retrives all checkpoints found in `path` directory.
+
+    Checkpoints are identified by matching filename to the specified pattern. If
+    the pattern contains groups, the result will be sorted by the first group in
+    descending order
+    """
+    pt_regexp = re.compile(pattern)
+    files = PathManager.ls(path)
+
+    entries = []
+    for i, f in enumerate(files):
+        m = pt_regexp.fullmatch(f)
+        if m is not None:
+            idx = float(m.group(1)) if len(m.groups()) > 0 else i
+            entries.append((idx, m.group(0)))
+    if keep_match:
+        return [(os.path.join(path, x[1]), x[0]) for x in sorted(entries, reverse=True)]
+    else:
+        return [os.path.join(path, x[1]) for x in sorted(entries, reverse=True)]
+
+
+def load_checkpoint_to_cpu(path, load_on_all_ranks=False):
+    """Loads a checkpoint to CPU (with upgrading for backward compatibility).
+
+    If doing single-GPU training or if the checkpoint is only being loaded by at
+    most one process on each node (current default behavior is for only rank 0
+    to read the checkpoint from disk), load_on_all_ranks should be False to
+    avoid errors from torch.distributed not having been initialized or
+    torch.distributed.barrier() hanging.
+
+    If all processes on each node may be loading the checkpoint
+    simultaneously, load_no_all_ranks should be set to True to avoid I/O
+    conflicts.
+
+    There's currently no support for > 1 but < all process loading the
+    checkpoint on each node.
+    """
+    local_path = PathManager.get_local_path(path)
+    # The locally cached file returned by get_local_path() may be stable for
+    # remote files that are periodically updated/overwritten (ex:
+    # checkpoint_last.pt) - so we remove the local copy, sync across processes
+    # (if needed), and then download a fresh copy.
+    if local_path != path and PathManager.path_requires_pathmanager(path):
+        try:
+            os.remove(local_path)
+        except FileNotFoundError:
+            # With potentially multiple processes removing the same file, the
+            # file being missing is benign (missing_ok isn't available until
+            # Python 3.8).
+            pass
+        if load_on_all_ranks:
+            torch.distributed.barrier()
+        local_path = PathManager.get_local_path(path)
+
+    with open(local_path, "rb") as f:
+        state = torch.load(f, map_location=torch.device("cpu"))
+
+    return state
+
+
+def verify_checkpoint_directory(save_dir: str) -> None:
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir, exist_ok=True)
+    temp_file_path = os.path.join(save_dir, "dummy")
+    try:
+        with open(temp_file_path, "w"):
+            pass
+    except OSError as e:
+        logger.warning("Unable to access checkpoint save directory: {}".format(save_dir))
+        raise e
+    else:
+        os.remove(temp_file_path)

genhpf/utils/data_utils.py

@@ -0,0 +1,130 @@
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+
+def compute_mask_indices(
+    shape: Tuple[int, int],
+    padding_mask: Optional[torch.Tensor],
+    mask_prob: float,
+    mask_length: int,
+    mask_type: str = "static",
+    mask_other: float = 0.0,
+    min_masks: int = 0,
+    no_overlap: bool = False,
+    min_space: int = 0,
+) -> np.ndarray:
+    """
+    Computes random mask spans for a given shape
+
+    Args:
+        shape: the the shape for which to compute masks.
+            should be of size 2 where first element is batch size and 2nd is timesteps
+        padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
+        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
+            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
+            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
+        mask_type: how to compute mask lengths
+            static = fixed size
+            uniform = sample from uniform distribution [mask_other, mask_length*2]
+            normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
+            poisson = sample from possion distribution with lambda = mask length
+        min_masks: minimum number of masked spans
+        no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
+        min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
+    """
+
+    bsz, all_sz = shape
+    mask = np.full((bsz, all_sz), False)
+
+    all_num_mask = int(
+        # add a random number for probabilistic rounding
+        mask_prob * all_sz / float(mask_length)
+        + np.random.rand()
+    )
+
+    all_num_mask = max(min_masks, all_num_mask)
+
+    mask_idcs = []
+    for i in range(bsz):
+        if padding_mask is not None:
+            sz = all_sz - padding_mask[i].long().sum().item()
+            num_mask = int(
+                # add a random number for probabilistic rounding
+                mask_prob * sz / float(mask_length)
+                + np.random.rand()
+            )
+            num_mask = max(min_masks, num_mask)
+        else:
+            sz = all_sz
+            num_mask = all_num_mask
+
+        if mask_type == "static":
+            lengths = np.full(num_mask, mask_length)
+        elif mask_type == "uniform":
+            lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask)
+        elif mask_type == "normal":
+            lengths = np.random.normal(mask_length, mask_other, size=num_mask)
+            lengths = [max(1, int(round(x))) for x in lengths]
+        elif mask_type == "poisson":
+            lengths = np.random.poisson(mask_length, size=num_mask)
+            lengths = [int(round(x)) for x in lengths]
+        else:
+            raise Exception("unknown mask selection " + mask_type)
+
+        if sum(lengths) == 0:
+            lengths[0] = min(mask_length, sz - 1)
+
+        if no_overlap:
+            mask_idc = []
+
+            def arrange(s, e, length, keep_length):
+                span_start = np.random.randint(s, e - length)
+                mask_idc.extend(span_start + i for i in range(length))
+
+                new_parts = []
+                if span_start - s - min_space >= keep_length:
+                    new_parts.append((s, span_start - min_space + 1))
+                if e - span_start - keep_length - min_space > keep_length:
+                    new_parts.append((span_start + length + min_space, e))
+                return new_parts
+
+            parts = [(0, sz)]
+            min_length = min(lengths)
+            for length in sorted(lengths, reverse=True):
+                lens = np.fromiter(
+                    (e - s if e - s >= length + min_space else 0 for s, e in parts),
+                    np.int,
+                )
+                l_sum = np.sum(lens)
+                if l_sum == 0:
+                    break
+                probs = lens / np.sum(lens)
+                c = np.random.choice(len(parts), p=probs)
+                s, e = parts.pop(c)
+                parts.extend(arrange(s, e, length, min_length))
+            mask_idc = np.asarray(mask_idc)
+        else:
+            min_len = min(lengths)
+            if sz - min_len <= num_mask:
+                min_len = sz - num_mask - 1
+
+            mask_idc = np.random.choice(sz - min_len, num_mask, replace=False)
+
+            mask_idc = np.asarray(
+                [
+                    mask_idc[j] + offset
+                    for j in range(len(mask_idc))
+                    for offset in range(lengths[j])
+                ]
+            )
+
+        mask_idcs.append(np.unique(mask_idc[mask_idc < sz]))
+
+    min_len = min([len(m) for m in mask_idcs])
+    for i, mask_idc in enumerate(mask_idcs):
+        if len(mask_idc) > min_len:
+            mask_idc = np.random.choice(mask_idc, min_len, replace=False)
+        mask[i, mask_idc] = True
+
+    return mask