reach-vb's picture
reach-vb HF staff
Stereo demo update (#60)
5325fcc
raw
history blame
14.3 kB
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from concurrent.futures import ThreadPoolExecutor
from collections import deque
from functools import partial
from hashlib import sha1
import logging
from pathlib import Path
import sys
import typing as tp
import zipfile
import flashy
import torch
logger = logging.getLogger(__name__)
def get_full_embed(full_embed: torch.Tensor, x: tp.Any, idx: int, device: tp.Union[str, torch.device]) -> torch.Tensor:
"""Utility function for the EmbeddingCache, returning the full embedding without any chunking.
This method can be used in case there is no need in extracting a chunk of the full embedding
read from the cache.
Args:
full_embed (torch.Tensor): The full embedding.
x (any): Batch object from which the full embedding is derived.
idx (torch.Tensor): Index of object to consider in the batch object.
Returns:
full_embed (torch.Tensor): The full embedding
"""
return full_embed.to(device)
class EmbeddingCache:
"""Cache around embeddings computation for faster execution.
The EmbeddingCache is storing pre-computed embeddings on disk and provides a simple API
to retrieve the pre-computed embeddings on full inputs and extract only a given chunk
using a user-provided function. When the cache is warm (all embeddings are pre-computed),
the EmbeddingCache allows for faster training as it removes the need of computing the embeddings.
Additionally, it provides in-memory cache around the loaded embeddings to limit IO footprint
and synchronization points in the forward calls.
Args:
cache_path (Path): Path to folder where all pre-computed embeddings are saved on disk.
device (str or torch.device): Device on which the embedding is returned.
compute_embed_fn (callable[[Path, any, int], torch.Tensor], optional): Function to compute
the embedding from a given object and path. This user provided function can compute the
embedding from the provided object or using the provided path as entry point. The last parameter
specify the index corresponding to the current embedding in the object that can represent batch metadata.
extract_embed_fn (callable[[torch.Tensor, any, int], torch.Tensor], optional): Function to extract
the desired embedding chunk from the full embedding loaded from the cache. The last parameter
specify the index corresponding to the current embedding in the object that can represent batch metadata.
If not specified, will return the full embedding unmodified.
"""
def __init__(self, cache_path: tp.Union[str, Path], device: tp.Union[str, torch.device],
compute_embed_fn: tp.Callable[[Path, tp.Any, int], torch.Tensor],
extract_embed_fn: tp.Optional[tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor]] = None):
self.cache_path = Path(cache_path)
self.device = device
self._compute_embed_fn = compute_embed_fn
self._extract_embed_fn: tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor]
if extract_embed_fn is not None:
self._extract_embed_fn = extract_embed_fn
else:
self._extract_embed_fn = partial(get_full_embed, device=device)
if self.cache_path is not None:
self.cache_path.mkdir(exist_ok=True, parents=True)
logger.info(f"Cache instantiated at: {self.cache_path}")
self.pool = ThreadPoolExecutor(8)
self.pool.__enter__()
self._current_batch_cache: dict = {}
self._memory_cache: dict = {}
def _get_cache_path(self, path: tp.Union[Path, str]):
"""Get cache path for the given file path."""
sig = sha1(str(path).encode()).hexdigest()
return self.cache_path / sig
@staticmethod
def _get_full_embed_from_cache(cache: Path):
"""Loads full pre-computed embedding from the cache."""
try:
embed = torch.load(cache, 'cpu')
except Exception as exc:
logger.error("Error loading %s: %r", cache, exc)
embed = None
return embed
def get_embed_from_cache(self, paths: tp.List[Path], x: tp.Any) -> torch.Tensor:
"""Get embedding from cache, computing and storing it to cache if not already cached.
The EmbeddingCache first tries to load the embedding from the in-memory cache
containing the pre-computed chunks populated through `populate_embed_cache`.
If not found, the full embedding is computed and stored on disk to be later accessed
to populate the in-memory cache, and the desired embedding chunk is extracted and returned.
Args:
paths (list[Path or str]): List of paths from where the embeddings can be loaded.
x (any): Object from which the embedding is extracted.
"""
embeds = []
for idx, path in enumerate(paths):
cache = self._get_cache_path(path)
if cache in self._current_batch_cache:
embed = self._current_batch_cache[cache]
else:
full_embed = self._compute_embed_fn(path, x, idx)
try:
with flashy.utils.write_and_rename(cache, pid=True) as f:
torch.save(full_embed.cpu(), f)
except Exception as exc:
logger.error('Error saving embed %s (%s): %r', cache, full_embed.shape, exc)
else:
logger.info('New embed cache saved: %s (%s)', cache, full_embed.shape)
embed = self._extract_embed_fn(full_embed, x, idx)
embeds.append(embed)
embed = torch.stack(embeds, dim=0)
return embed
def populate_embed_cache(self, paths: tp.List[Path], x: tp.Any) -> None:
"""Populate in-memory caches for embeddings reading from the embeddings stored on disk.
The in-memory caches consist in a cache for the full embedding and another cache for the
final embedding chunk. Such caches are used to limit the IO access when computing the actual embeddings
and reduce the IO footprint and synchronization points during forward passes.
Args:
paths (list[Path]): List of paths from where the embeddings can be loaded.
x (any): Object from which the embedding is extracted.
"""
self._current_batch_cache.clear()
if self.cache_path is not None:
futures: list = []
for path in paths:
assert path is not None, "Path is required for computation from cache"
cache = self._get_cache_path(path)
if cache in self._memory_cache or not cache.exists():
futures.append(None)
else:
futures.append(self.pool.submit(EmbeddingCache._get_full_embed_from_cache, cache))
for idx, (path, future) in enumerate(zip(paths, futures)):
assert path is not None
cache = self._get_cache_path(path)
full_embed = None
if future is None:
if cache in self._memory_cache:
full_embed = self._memory_cache[cache]
else:
full_embed = future.result()
if full_embed is not None:
self._memory_cache[cache] = full_embed
full_embed = full_embed.to(self.device)
if full_embed is not None:
embed = self._extract_embed_fn(full_embed, x, idx)
self._current_batch_cache[cache] = embed
class CachedBatchWriter:
"""Write pre computed caches for mini batches. This can
make loading a lot more efficient depending on your filesystem.
Args:
cache_folder (Path): folder in which the cached minibatches
will be stored.
Inside cache folder, the structure is the following:
`epoch_number / update_number.zip`
And the zip file contains one entry per batch item.
It is possible to use the cache with a batch size smaller than
created with but obviously not larger. Make sure to call the
`start_epoch(epoch)` method for indicating changes of epochs.
See the grid `audiocraft/grids/musicgen/musicgen_warmup_cache.py`
for an example of how to warmup the cache.
"""
def __init__(self, cache_folder: Path):
self.cache_folder = cache_folder
self._current_epoch: tp.Optional[int] = None
self._current_index = 0
def start_epoch(self, epoch: int):
"""Call at the beginning of each epoch.
"""
self._current_epoch = epoch
self._current_index = 0
self._zip_path.parent.mkdir(exist_ok=True, parents=True)
@staticmethod
def _get_zip_path(cache_folder: Path, epoch: int, index: int):
return cache_folder / f"{epoch:05d}" / f"{index:06d}.zip"
@property
def _zip_path(self):
assert self._current_epoch is not None
return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, self._current_index)
def save(self, *content):
"""Save one mini batch. This function is distributed-aware
and will automatically merge all the items from the different
workers.
"""
all_contents = []
for rank in range(flashy.distrib.world_size()):
their_content = flashy.distrib.broadcast_object(content, src=rank)
all_contents.append(their_content)
if flashy.distrib.is_rank_zero():
idx = 0
with flashy.utils.write_and_rename(self._zip_path) as tmp:
with zipfile.ZipFile(tmp, 'w') as zf:
for content in all_contents:
for vals in zip(*content):
with zf.open(f'{idx}', 'w') as f: # type: ignore
torch.save(vals, f)
idx += 1
flashy.distrib.barrier()
self._current_index += 1
class CachedBatchLoader:
"""Loader for cached mini-batches dumped with `CachedBatchWriter`.
Args:
cache_folder (Path): folder in which the cached minibatches are stored.
batch_size (int): batch size (per GPU) expected.
num_workers (int): number of workers to use for loading.
min_length (int): minimum expected length for each epoch. If some
mini-batches are missing, and error is raised.
This is iterable just like a regular DataLoader.
"""
def __init__(self, cache_folder: Path, batch_size: int,
num_workers: int = 10, min_length: int = 1):
self.cache_folder = cache_folder
self.batch_size = batch_size
self.num_workers = num_workers
self.min_length = min_length
self._current_epoch: tp.Optional[int] = None
self.sampler = None # for compatibility with the regular DataLoader
def __len__(self):
path = CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch or 0, 0).parent
return len([p for p in path.iterdir() if p.suffix == ".zip"])
def start_epoch(self, epoch: int):
"""Call at the beginning of each epoch.
"""
self._current_epoch = epoch
def _zip_path(self, index: int):
assert self._current_epoch is not None
return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, index)
def _load_one(self, index: int):
zip_path = self._zip_path(index)
if not zip_path.exists():
if index < self.min_length:
raise RuntimeError(f"Cache should have at least {self.min_length} batches, but {index} doesn't exist")
return None
mode = "rb" if sys.version_info >= (3, 9) else "r"
try:
with zipfile.ZipFile(zip_path, 'r') as zf:
rank = flashy.distrib.rank()
world_size = flashy.distrib.world_size()
root = zipfile.Path(zf)
items = list(root.iterdir())
total_batch_size = self.batch_size * world_size
if len(items) < total_batch_size:
raise RuntimeError(
f"The cache can handle a max batch size of {len(items)}, "
f"but {total_batch_size} is needed.")
start = rank * self.batch_size
items = items[start: start + self.batch_size]
assert len(items) == self.batch_size
entries = []
entries = [torch.load(item.open(mode), 'cpu') for item in items] # type: ignore
transposed = zip(*entries)
out = []
for part in transposed:
assert len(part) > 0
if isinstance(part[0], torch.Tensor):
out.append(torch.stack(part))
else:
out.append(part)
return out
except Exception:
logger.error("Error when reading zip path %s", zip_path)
raise
def __iter__(self):
"""This will yields tuples, exactly as provided to the
`CachedBatchWriter.save` method.
"""
pool = ThreadPoolExecutor(self.num_workers)
next_index = 0
queue = deque()
def _get_next():
nonlocal next_index
r = queue.popleft().result()
if r is None:
return None
else:
queue.append(pool.submit(self._load_one, next_index))
next_index += 1
return r
with pool:
# fill the buffer of fetching jobs.
for _ in range(2 * self.num_workers):
queue.append(pool.submit(self._load_one, next_index))
next_index += 1
while True:
batch = _get_next()
if batch is None:
return
yield batch