xinference 1.10.0__py3-none-any.whl → 1.10.1__py3-none-any.whl

xinference/thirdparty/audiotools/ml/accelerator.py

@@ -0,0 +1,184 @@
+import os
+import typing
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DataParallel
+from torch.nn.parallel import DistributedDataParallel
+
+from ..data.datasets import ResumableDistributedSampler as DistributedSampler
+from ..data.datasets import ResumableSequentialSampler as SequentialSampler
+
+
+class Accelerator:  # pragma: no cover
+    """This class is used to prepare models and dataloaders for
+    usage with DDP or DP. Use the functions prepare_model, prepare_dataloader to
+    prepare the respective objects. In the case of models, they are moved to
+    the appropriate GPU and SyncBatchNorm is applied to them. In the case of
+    dataloaders, a sampler is created and the dataloader is initialized with
+    that sampler.
+
+    If the world size is 1, prepare_model and prepare_dataloader are
+    no-ops. If the environment variable ``LOCAL_RANK`` is not set, then the
+    script was launched without ``torchrun``, and ``DataParallel``
+    will be used instead of ``DistributedDataParallel`` (not recommended), if
+    the world size (number of GPUs) is greater than 1.
+
+    Parameters
+    ----------
+    amp : bool, optional
+        Whether or not to enable automatic mixed precision, by default False
+    """
+
+    def __init__(self, amp: bool = False):
+        local_rank = os.getenv("LOCAL_RANK", None)
+        self.world_size = torch.cuda.device_count()
+
+        self.use_ddp = self.world_size > 1 and local_rank is not None
+        self.use_dp = self.world_size > 1 and local_rank is None
+        self.device = "cpu" if self.world_size == 0 else "cuda"
+
+        if self.use_ddp:
+            local_rank = int(local_rank)
+            dist.init_process_group(
+                "nccl",
+                init_method="env://",
+                world_size=self.world_size,
+                rank=local_rank,
+            )
+
+        self.local_rank = 0 if local_rank is None else local_rank
+        self.amp = amp
+
+        class DummyScaler:
+            def __init__(self):
+                pass
+
+            def step(self, optimizer):
+                optimizer.step()
+
+            def scale(self, loss):
+                return loss
+
+            def unscale_(self, optimizer):
+                return optimizer
+
+            def update(self):
+                pass
+
+        self.scaler = torch.cuda.amp.GradScaler() if amp else DummyScaler()
+        self.device_ctx = (
+            torch.cuda.device(self.local_rank) if torch.cuda.is_available() else None
+        )
+
+    def __enter__(self):
+        if self.device_ctx is not None:
+            self.device_ctx.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.device_ctx is not None:
+            self.device_ctx.__exit__(exc_type, exc_value, traceback)
+
+    def prepare_model(self, model: torch.nn.Module, **kwargs):
+        """Prepares model for DDP or DP. The model is moved to
+        the device of the correct rank.
+
+        Parameters
+        ----------
+        model : torch.nn.Module
+            Model that is converted for DDP or DP.
+
+        Returns
+        -------
+        torch.nn.Module
+            Wrapped model, or original model if DDP and DP are turned off.
+        """
+        model = model.to(self.device)
+        if self.use_ddp:
+            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+            model = DistributedDataParallel(
+                model, device_ids=[self.local_rank], **kwargs
+            )
+        elif self.use_dp:
+            model = DataParallel(model, **kwargs)
+        return model
+
+    # Automatic mixed-precision utilities
+    def autocast(self, *args, **kwargs):
+        """Context manager for autocasting. Arguments
+        go to ``torch.cuda.amp.autocast``.
+        """
+        return torch.cuda.amp.autocast(self.amp, *args, **kwargs)
+
+    def backward(self, loss: torch.Tensor):
+        """Backwards pass, after scaling the loss if ``amp`` is
+        enabled.
+
+        Parameters
+        ----------
+        loss : torch.Tensor
+            Loss value.
+        """
+        self.scaler.scale(loss).backward()
+
+    def step(self, optimizer: torch.optim.Optimizer):
+        """Steps the optimizer, using a ``scaler`` if ``amp`` is
+        enabled.
+
+        Parameters
+        ----------
+        optimizer : torch.optim.Optimizer
+            Optimizer to step forward.
+        """
+        self.scaler.step(optimizer)
+
+    def update(self):
+        """Updates the scale factor."""
+        self.scaler.update()
+
+    def prepare_dataloader(
+        self, dataset: typing.Iterable, start_idx: int = None, **kwargs
+    ):
+        """Wraps a dataset with a DataLoader, using the correct sampler if DDP is
+        enabled.
+
+        Parameters
+        ----------
+        dataset : typing.Iterable
+            Dataset to build Dataloader around.
+        start_idx : int, optional
+            Start index of sampler, useful if resuming from some epoch,
+            by default None
+
+        Returns
+        -------
+        _type_
+            _description_
+        """
+
+        if self.use_ddp:
+            sampler = DistributedSampler(
+                dataset,
+                start_idx,
+                num_replicas=self.world_size,
+                rank=self.local_rank,
+            )
+            if "num_workers" in kwargs:
+                kwargs["num_workers"] = max(kwargs["num_workers"] // self.world_size, 1)
+            kwargs["batch_size"] = max(kwargs["batch_size"] // self.world_size, 1)
+        else:
+            sampler = SequentialSampler(dataset, start_idx)
+
+        dataloader = torch.utils.data.DataLoader(dataset, sampler=sampler, **kwargs)
+        return dataloader
+
+    @staticmethod
+    def unwrap(model):
+        """Unwraps the model if it was wrapped in DDP or DP, otherwise
+        just returns the model. Use this to unwrap the model returned by
+        :py:func:`audiotools.ml.accelerator.Accelerator.prepare_model`.
+        """
+        if hasattr(model, "module"):
+            return model.module
+        return model

xinference/thirdparty/audiotools/ml/decorators.py

@@ -0,0 +1,440 @@
+import math
+import os
+import time
+from collections import defaultdict
+from functools import wraps
+
+import torch
+import torch.distributed as dist
+from rich import box
+from rich.console import Console
+from rich.console import Group
+from rich.live import Live
+from rich.markdown import Markdown
+from rich.padding import Padding
+from rich.panel import Panel
+from rich.progress import BarColumn
+from rich.progress import Progress
+from rich.progress import SpinnerColumn
+from rich.progress import TimeElapsedColumn
+from rich.progress import TimeRemainingColumn
+from rich.rule import Rule
+from rich.table import Table
+from torch.utils.tensorboard import SummaryWriter
+
+
+# This is here so that the history can be pickled.
+def default_list():
+    return []
+
+
+class Mean:
+    """Keeps track of the running mean, along with the latest
+    value.
+    """
+
+    def __init__(self):
+        self.reset()
+
+    def __call__(self):
+        mean = self.total / max(self.count, 1)
+        return mean
+
+    def reset(self):
+        self.count = 0
+        self.total = 0
+
+    def update(self, val):
+        if math.isfinite(val):
+            self.count += 1
+            self.total += val
+
+
+def when(condition):
+    """Runs a function only when the condition is met. The condition is
+    a function that is run.
+
+    Parameters
+    ----------
+    condition : Callable
+        Function to run to check whether or not to run the decorated
+        function.
+
+    Example
+    -------
+    Checkpoint only runs every 100 iterations, and only if the
+    local rank is 0.
+
+    >>> i = 0
+    >>> rank = 0
+    >>>
+    >>> @when(lambda: i % 100 == 0 and rank == 0)
+    >>> def checkpoint():
+    >>>     print("Saving to /runs/exp1")
+    >>>
+    >>> for i in range(1000):
+    >>>     checkpoint()
+
+    """
+
+    def decorator(fn):
+        @wraps(fn)
+        def decorated(*args, **kwargs):
+            if condition():
+                return fn(*args, **kwargs)
+
+        return decorated
+
+    return decorator
+
+
+def timer(prefix: str = "time"):
+    """Adds execution time to the output dictionary of the decorated
+    function. The function decorated by this must output a dictionary.
+    The key added will follow the form "[prefix]/[name_of_function]"
+
+    Parameters
+    ----------
+    prefix : str, optional
+        The key added will follow the form "[prefix]/[name_of_function]",
+        by default "time".
+    """
+
+    def decorator(fn):
+        @wraps(fn)
+        def decorated(*args, **kwargs):
+            s = time.perf_counter()
+            output = fn(*args, **kwargs)
+            assert isinstance(output, dict)
+            e = time.perf_counter()
+            output[f"{prefix}/{fn.__name__}"] = e - s
+            return output
+
+        return decorated
+
+    return decorator
+
+
+class Tracker:
+    """
+    A tracker class that helps to monitor the progress of training and logging the metrics.
+
+    Attributes
+    ----------
+    metrics : dict
+        A dictionary containing the metrics for each label.
+    history : dict
+        A dictionary containing the history of metrics for each label.
+    writer : SummaryWriter
+        A SummaryWriter object for logging the metrics.
+    rank : int
+        The rank of the current process.
+    step : int
+        The current step of the training.
+    tasks : dict
+        A dictionary containing the progress bars and tables for each label.
+    pbar : Progress
+        A progress bar object for displaying the progress.
+    consoles : list
+        A list of console objects for logging.
+    live : Live
+        A Live object for updating the display live.
+
+    Methods
+    -------
+    print(msg: str)
+        Prints the given message to all consoles.
+    update(label: str, fn_name: str)
+        Updates the progress bar and table for the given label.
+    done(label: str, title: str)
+        Resets the progress bar and table for the given label and prints the final result.
+    track(label: str, length: int, completed: int = 0, op: dist.ReduceOp = dist.ReduceOp.AVG, ddp_active: bool = "LOCAL_RANK" in os.environ)
+        A decorator for tracking the progress and metrics of a function.
+    log(label: str, value_type: str = "value", history: bool = True)
+        A decorator for logging the metrics of a function.
+    is_best(label: str, key: str) -> bool
+        Checks if the latest value of the given key in the label is the best so far.
+    state_dict() -> dict
+        Returns a dictionary containing the state of the tracker.
+    load_state_dict(state_dict: dict) -> Tracker
+        Loads the state of the tracker from the given state dictionary.
+    """
+
+    def __init__(
+        self,
+        writer: SummaryWriter = None,
+        log_file: str = None,
+        rank: int = 0,
+        console_width: int = 100,
+        step: int = 0,
+    ):
+        """
+        Initializes the Tracker object.
+
+        Parameters
+        ----------
+        writer : SummaryWriter, optional
+            A SummaryWriter object for logging the metrics, by default None.
+        log_file : str, optional
+            The path to the log file, by default None.
+        rank : int, optional
+            The rank of the current process, by default 0.
+        console_width : int, optional
+            The width of the console, by default 100.
+        step : int, optional
+            The current step of the training, by default 0.
+        """
+        self.metrics = {}
+        self.history = {}
+        self.writer = writer
+        self.rank = rank
+        self.step = step
+
+        # Create progress bars etc.
+        self.tasks = {}
+        self.pbar = Progress(
+            SpinnerColumn(),
+            "[progress.description]{task.description}",
+            "{task.completed}/{task.total}",
+            BarColumn(),
+            TimeElapsedColumn(),
+            "/",
+            TimeRemainingColumn(),
+        )
+        self.consoles = [Console(width=console_width)]
+        self.live = Live(console=self.consoles[0], refresh_per_second=10)
+        if log_file is not None:
+            self.consoles.append(Console(width=console_width, file=open(log_file, "a")))
+
+    def print(self, msg):
+        """
+        Prints the given message to all consoles.
+
+        Parameters
+        ----------
+        msg : str
+            The message to be printed.
+        """
+        if self.rank == 0:
+            for c in self.consoles:
+                c.log(msg)
+
+    def update(self, label, fn_name):
+        """
+        Updates the progress bar and table for the given label.
+
+        Parameters
+        ----------
+        label : str
+            The label of the progress bar and table to be updated.
+        fn_name : str
+            The name of the function associated with the label.
+        """
+        if self.rank == 0:
+            self.pbar.advance(self.tasks[label]["pbar"])
+
+            # Create table
+            table = Table(title=label, expand=True, box=box.MINIMAL)
+            table.add_column("key", style="cyan")
+            table.add_column("value", style="bright_blue")
+            table.add_column("mean", style="bright_green")
+
+            keys = self.metrics[label]["value"].keys()
+            for k in keys:
+                value = self.metrics[label]["value"][k]
+                mean = self.metrics[label]["mean"][k]()
+                table.add_row(k, f"{value:10.6f}", f"{mean:10.6f}")
+
+            self.tasks[label]["table"] = table
+            tables = [t["table"] for t in self.tasks.values()]
+            group = Group(*tables, self.pbar)
+            self.live.update(
+                Group(
+                    Padding("", (0, 0)),
+                    Rule(f"[italic]{fn_name}()", style="white"),
+                    Padding("", (0, 0)),
+                    Panel.fit(
+                        group, padding=(0, 5), title="[b]Progress", border_style="blue"
+                    ),
+                )
+            )
+
+    def done(self, label: str, title: str):
+        """
+        Resets the progress bar and table for the given label and prints the final result.
+
+        Parameters
+        ----------
+        label : str
+            The label of the progress bar and table to be reset.
+        title : str
+            The title to be displayed when printing the final result.
+        """
+        for label in self.metrics:
+            for v in self.metrics[label]["mean"].values():
+                v.reset()
+
+        if self.rank == 0:
+            self.pbar.reset(self.tasks[label]["pbar"])
+            tables = [t["table"] for t in self.tasks.values()]
+            group = Group(Markdown(f"# {title}"), *tables, self.pbar)
+            self.print(group)
+
+    def track(
+        self,
+        label: str,
+        length: int,
+        completed: int = 0,
+        op: dist.ReduceOp = dist.ReduceOp.AVG,
+        ddp_active: bool = "LOCAL_RANK" in os.environ,
+    ):
+        """
+        A decorator for tracking the progress and metrics of a function.
+
+        Parameters
+        ----------
+        label : str
+            The label to be associated with the progress and metrics.
+        length : int
+            The total number of iterations to be completed.
+        completed : int, optional
+            The number of iterations already completed, by default 0.
+        op : dist.ReduceOp, optional
+            The reduce operation to be used, by default dist.ReduceOp.AVG.
+        ddp_active : bool, optional
+            Whether the DistributedDataParallel is active, by default "LOCAL_RANK" in os.environ.
+        """
+        self.tasks[label] = {
+            "pbar": self.pbar.add_task(
+                f"[white]Iteration ({label})", total=length, completed=completed
+            ),
+            "table": Table(),
+        }
+        self.metrics[label] = {
+            "value": defaultdict(),
+            "mean": defaultdict(lambda: Mean()),
+        }
+
+        def decorator(fn):
+            @wraps(fn)
+            def decorated(*args, **kwargs):
+                output = fn(*args, **kwargs)
+                if not isinstance(output, dict):
+                    self.update(label, fn.__name__)
+                    return output
+                # Collect across all DDP processes
+                scalar_keys = []
+                for k, v in output.items():
+                    if isinstance(v, (int, float)):
+                        v = torch.tensor([v])
+                    if not torch.is_tensor(v):
+                        continue
+                    if ddp_active and v.is_cuda:  # pragma: no cover
+                        dist.all_reduce(v, op=op)
+                    output[k] = v.detach()
+                    if torch.numel(v) == 1:
+                        scalar_keys.append(k)
+                        output[k] = v.item()
+
+                # Save the outputs to tracker
+                for k, v in output.items():
+                    if k not in scalar_keys:
+                        continue
+                    self.metrics[label]["value"][k] = v
+                    # Update the running mean
+                    self.metrics[label]["mean"][k].update(v)
+
+                self.update(label, fn.__name__)
+                return output
+
+            return decorated
+
+        return decorator
+
+    def log(self, label: str, value_type: str = "value", history: bool = True):
+        """
+        A decorator for logging the metrics of a function.
+
+        Parameters
+        ----------
+        label : str
+            The label to be associated with the logging.
+        value_type : str, optional
+            The type of value to be logged, by default "value".
+        history : bool, optional
+            Whether to save the history of the metrics, by default True.
+        """
+        assert value_type in ["mean", "value"]
+        if history:
+            if label not in self.history:
+                self.history[label] = defaultdict(default_list)
+
+        def decorator(fn):
+            @wraps(fn)
+            def decorated(*args, **kwargs):
+                output = fn(*args, **kwargs)
+                if self.rank == 0:
+                    nonlocal value_type, label
+                    metrics = self.metrics[label][value_type]
+                    for k, v in metrics.items():
+                        v = v() if isinstance(v, Mean) else v
+                        if self.writer is not None:
+                            self.writer.add_scalar(f"{k}/{label}", v, self.step)
+                        if label in self.history:
+                            self.history[label][k].append(v)
+
+                    if label in self.history:
+                        self.history[label]["step"].append(self.step)
+
+                return output
+
+            return decorated
+
+        return decorator
+
+    def is_best(self, label, key):
+        """
+        Checks if the latest value of the given key in the label is the best so far.
+
+        Parameters
+        ----------
+        label : str
+            The label of the metrics to be checked.
+        key : str
+            The key of the metric to be checked.
+
+        Returns
+        -------
+        bool
+            True if the latest value is the best so far, otherwise False.
+        """
+        return self.history[label][key][-1] == min(self.history[label][key])
+
+    def state_dict(self):
+        """
+        Returns a dictionary containing the state of the tracker.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the history and step of the tracker.
+        """
+        return {"history": self.history, "step": self.step}
+
+    def load_state_dict(self, state_dict):
+        """
+        Loads the state of the tracker from the given state dictionary.
+
+        Parameters
+        ----------
+        state_dict : dict
+            A dictionary containing the history and step of the tracker.
+
+        Returns
+        -------
+        Tracker
+            The tracker object with the loaded state.
+        """
+        self.history = state_dict["history"]
+        self.step = state_dict["step"]
+        return self

xinference/thirdparty/audiotools/ml/experiment.py

@@ -0,0 +1,90 @@
+"""
+Useful class for Experiment tracking, and ensuring code is
+saved alongside files.
+"""  # fmt: skip
+import datetime
+import os
+import shlex
+import shutil
+import subprocess
+import typing
+from pathlib import Path
+
+import randomname
+
+
+class Experiment:
+    """This class contains utilities for managing experiments.
+    It is a context manager, that when you enter it, changes
+    your directory to a specified experiment folder (which
+    optionally can have an automatically generated experiment
+    name, or a specified one), and changes the CUDA device used
+    to the specified device (or devices).
+
+    Parameters
+    ----------
+    exp_directory : str
+        Folder where all experiments are saved, by default "runs/".
+    exp_name : str, optional
+        Name of the experiment, by default uses the current time, date, and
+        hostname to save.
+    """
+
+    def __init__(
+        self,
+        exp_directory: str = "runs/",
+        exp_name: str = None,
+    ):
+        if exp_name is None:
+            exp_name = self.generate_exp_name()
+        exp_dir = Path(exp_directory) / exp_name
+        exp_dir.mkdir(parents=True, exist_ok=True)
+
+        self.exp_dir = exp_dir
+        self.exp_name = exp_name
+        self.git_tracked_files = (
+            subprocess.check_output(
+                shlex.split("git ls-tree --full-tree --name-only -r HEAD")
+            )
+            .decode("utf-8")
+            .splitlines()
+        )
+        self.parent_directory = Path(".").absolute()
+
+    def __enter__(self):
+        self.prev_dir = os.getcwd()
+        os.chdir(self.exp_dir)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        os.chdir(self.prev_dir)
+
+    @staticmethod
+    def generate_exp_name():
+        """Generates a random experiment name based on the date
+        and a randomly generated adjective-noun tuple.
+
+        Returns
+        -------
+        str
+            Randomly generated experiment name.
+        """
+        date = datetime.datetime.now().strftime("%y%m%d")
+        name = f"{date}-{randomname.get_name()}"
+        return name
+
+    def snapshot(self, filter_fn: typing.Callable = lambda f: True):
+        """Captures a full snapshot of all the files tracked by git at the time
+        the experiment is run. It also captures the diff against the committed
+        code as a separate file.
+
+        Parameters
+        ----------
+        filter_fn : typing.Callable, optional
+            Function that can be used to exclude some files
+            from the snapshot, by default accepts all files
+        """
+        for f in self.git_tracked_files:
+            if filter_fn(f):
+                Path(f).parent.mkdir(parents=True, exist_ok=True)
+                shutil.copyfile(self.parent_directory / f, f)