webdataset.shardlists
Train PyTorch models directly from POSIX tar archive.
Code works locally or over HTTP connections.
View Source
# # Copyright (c) 2017-2021 NVIDIA CORPORATION. All rights reserved. # This file is part of the WebDataset library. # See the LICENSE file for licensing terms (BSD-style). # """Train PyTorch models directly from POSIX tar archive. Code works locally or over HTTP connections. """ import os, random, sys, time from dataclasses import dataclass, field from itertools import islice from typing import List import braceexpand, yaml from . import utils from .filters import pipelinefilter from .pytorch import IterableDataset def expand_urls(urls): if isinstance(urls, str): urllist = urls.split("::") result = [] for url in urllist: result.extend(braceexpand.braceexpand(url)) return result else: return list(urls) class SimpleShardList(IterableDataset): """An iterable dataset yielding a list of urls.""" def __init__(self, urls, seed=None): """Iterate through the list of shards. :param urls: a list of URLs as a Python list or brace notation string """ super().__init__() urls = expand_urls(urls) self.urls = urls assert isinstance(self.urls[0], str) self.seed = seed def __len__(self): return len(self.urls) def __iter__(self): """Return an iterator over the shards.""" urls = self.urls.copy() if self.seed is not None: random.Random(self.seed).shuffle(urls) for url in urls: yield dict(url=url) def split_by_node(src, group=None): rank, world_size, worker, num_workers = utils.pytorch_worker_info(group=group) if world_size > 1: for s in islice(src, rank, None, world_size): yield s else: for s in src: yield s def single_node_only(src, group=None): rank, world_size, worker, num_workers = utils.pytorch_worker_info(group=group) if world_size > 1: raise ValueError( "input pipeline needs to be reconfigured for multinode training" ) for s in src: yield s def split_by_worker(src): rank, world_size, worker, num_workers = utils.pytorch_worker_info() if num_workers > 1: for s in islice(src, worker, None, num_workers): yield s else: for s in src: yield s def resampled_(src, n=sys.maxsize): import random seed = time.time_ns() try: seed = open("/dev/random", "rb").read(20) except Exception as exn: print(repr(exn)[:50], file=sys.stderr) rng = random.Random(seed) print("# resampled loading", file=sys.stderr) items = list(src) print(f"# resampled got {len(items)} samples, yielding {n}", file=sys.stderr) for i in range(n): yield rng.choice(items) resampled = pipelinefilter(resampled_) def non_empty(src): count = 0 for s in src: yield s count += 1 if count == 0: raise ValueError( "pipeline stage received no data at all and this was declared as an error" ) @dataclass class MSSource: """Class representing a data source.""" name: str = "" perepoch: int = -1 resample: bool = False urls: List[str] = field(default_factory=list) default_rng = random.Random() def expand(s): return os.path.expanduser(os.path.expandvars(s)) class MultiShardSample(IterableDataset): def __init__(self, fname): """Construct a shardlist from multiple sources using a YAML spec.""" self.epoch = -1 self.parse_spec(fname) def parse_spec(self, fname): self.rng = default_rng # capture default_rng if we fork if isinstance(fname, dict): spec = fname fname = "{dict}" else: with open(fname) as stream: spec = yaml.safe_load(stream) assert set(spec.keys()).issubset(set("prefix datasets buckets".split())), list( spec.keys() ) prefix = expand(spec.get("prefix", "")) self.sources = [] for ds in spec["datasets"]: assert set(ds.keys()).issubset( set("buckets name shards resample choose".split()) ), list(ds.keys()) buckets = ds.get("buckets", spec.get("buckets", [])) if isinstance(buckets, str): buckets = [buckets] buckets = [expand(s) for s in buckets] if buckets == []: buckets = [""] assert ( len(buckets) == 1 ), f"{buckets}: FIXME support for multiple buckets unimplemented" bucket = buckets[0] name = ds.get("name", "@" + bucket) urls = ds["shards"] if isinstance(urls, str): urls = [urls] # urls = [u for url in urls for u in braceexpand.braceexpand(url)] urls = [ prefix + os.path.join(bucket, u) for url in urls for u in braceexpand.braceexpand(expand(url)) ] resample = ds.get("resample", -1) nsample = ds.get("choose", -1) if nsample > len(urls): raise ValueError( f"perepoch {nsample} must be no greater than the number of shards" ) if (nsample > 0) and (resample > 0): raise ValueError("specify only one of perepoch or choose") entry = MSSource(name=name, urls=urls, perepoch=nsample, resample=resample) self.sources.append(entry) print(f"# {name} {len(urls)} {nsample}", file=sys.stderr) def set_epoch(self, seed): """Set the current epoch (for consistent shard selection among nodes).""" self.rng = random.Random(seed) def get_shards_for_epoch(self): result = [] for source in self.sources: if source.resample > 0: # sample with replacement l = self.rng.choices(source.urls, k=source.resample) elif source.perepoch > 0: # sample without replacement l = list(source.urls) self.rng.shuffle(l) l = l[: source.perepoch] else: l = list(source.urls) result += l self.rng.shuffle(result) return result def __iter__(self): shards = self.get_shards_for_epoch() for shard in shards: yield dict(url=shard) def shardspec(spec): if spec.endswith(".yaml"): return MultiShardSample(spec) else: return SimpleShardList(spec) class ResampledShards(IterableDataset): """An iterable dataset yielding a list of urls.""" def __init__( self, urls, nshards=sys.maxsize, worker_seed=None, deterministic=False, ): """Sample shards from the shard list with replacement. :param urls: a list of URLs as a Python list or brace notation string """ super().__init__() urls = expand_urls(urls) self.urls = urls assert isinstance(self.urls[0], str) self.nshards = nshards self.rng = random.Random() self.worker_seed = ( utils.pytorch_worker_seed if worker_seed is None else worker_seed ) self.deterministic = deterministic self.epoch = -1 def __iter__(self): """Return an iterator over the shards.""" self.epoch += 1 if self.deterministic: seed = (self.worker_seed(), self.epoch) else: seed = (self.worker_seed(), self.epoch, os.getpid(), time.time()) self.rng.seed(seed) for _ in range(self.nshards): yield dict(url=self.rng.choice(self.urls))
View Source
def expand_urls(urls): if isinstance(urls, str): urllist = urls.split("::") result = [] for url in urllist: result.extend(braceexpand.braceexpand(url)) return result else: return list(urls)
View Source
class SimpleShardList(IterableDataset): """An iterable dataset yielding a list of urls.""" def __init__(self, urls, seed=None): """Iterate through the list of shards. :param urls: a list of URLs as a Python list or brace notation string """ super().__init__() urls = expand_urls(urls) self.urls = urls assert isinstance(self.urls[0], str) self.seed = seed def __len__(self): return len(self.urls) def __iter__(self): """Return an iterator over the shards.""" urls = self.urls.copy() if self.seed is not None: random.Random(self.seed).shuffle(urls) for url in urls: yield dict(url=url)
An iterable dataset yielding a list of urls.
View Source
def __init__(self, urls, seed=None): """Iterate through the list of shards. :param urls: a list of URLs as a Python list or brace notation string """ super().__init__() urls = expand_urls(urls) self.urls = urls assert isinstance(self.urls[0], str) self.seed = seed
Iterate through the list of shards.
:param urls: a list of URLs as a Python list or brace notation string
Inherited Members
- torch.utils.data.dataset.IterableDataset
- functions
- reduce_ex_hook
- register_function
- register_datapipe_as_function
- set_reduce_ex_hook
- type
View Source
def split_by_node(src, group=None): rank, world_size, worker, num_workers = utils.pytorch_worker_info(group=group) if world_size > 1: for s in islice(src, rank, None, world_size): yield s else: for s in src: yield s
View Source
def single_node_only(src, group=None): rank, world_size, worker, num_workers = utils.pytorch_worker_info(group=group) if world_size > 1: raise ValueError( "input pipeline needs to be reconfigured for multinode training" ) for s in src: yield s
View Source
def split_by_worker(src): rank, world_size, worker, num_workers = utils.pytorch_worker_info() if num_workers > 1: for s in islice(src, worker, None, num_workers): yield s else: for s in src: yield s
View Source
def resampled_(src, n=sys.maxsize): import random seed = time.time_ns() try: seed = open("/dev/random", "rb").read(20) except Exception as exn: print(repr(exn)[:50], file=sys.stderr) rng = random.Random(seed) print("# resampled loading", file=sys.stderr) items = list(src) print(f"# resampled got {len(items)} samples, yielding {n}", file=sys.stderr) for i in range(n): yield rng.choice(items)
View Source
def non_empty(src): count = 0 for s in src: yield s count += 1 if count == 0: raise ValueError( "pipeline stage received no data at all and this was declared as an error" )
View Source
class MSSource: """Class representing a data source.""" name: str = "" perepoch: int = -1 resample: bool = False urls: List[str] = field(default_factory=list)
Class representing a data source.
View Source
def expand(s): return os.path.expanduser(os.path.expandvars(s))
View Source
class MultiShardSample(IterableDataset): def __init__(self, fname): """Construct a shardlist from multiple sources using a YAML spec.""" self.epoch = -1 self.parse_spec(fname) def parse_spec(self, fname): self.rng = default_rng # capture default_rng if we fork if isinstance(fname, dict): spec = fname fname = "{dict}" else: with open(fname) as stream: spec = yaml.safe_load(stream) assert set(spec.keys()).issubset(set("prefix datasets buckets".split())), list( spec.keys() ) prefix = expand(spec.get("prefix", "")) self.sources = [] for ds in spec["datasets"]: assert set(ds.keys()).issubset( set("buckets name shards resample choose".split()) ), list(ds.keys()) buckets = ds.get("buckets", spec.get("buckets", [])) if isinstance(buckets, str): buckets = [buckets] buckets = [expand(s) for s in buckets] if buckets == []: buckets = [""] assert ( len(buckets) == 1 ), f"{buckets}: FIXME support for multiple buckets unimplemented" bucket = buckets[0] name = ds.get("name", "@" + bucket) urls = ds["shards"] if isinstance(urls, str): urls = [urls] # urls = [u for url in urls for u in braceexpand.braceexpand(url)] urls = [ prefix + os.path.join(bucket, u) for url in urls for u in braceexpand.braceexpand(expand(url)) ] resample = ds.get("resample", -1) nsample = ds.get("choose", -1) if nsample > len(urls): raise ValueError( f"perepoch {nsample} must be no greater than the number of shards" ) if (nsample > 0) and (resample > 0): raise ValueError("specify only one of perepoch or choose") entry = MSSource(name=name, urls=urls, perepoch=nsample, resample=resample) self.sources.append(entry) print(f"# {name} {len(urls)} {nsample}", file=sys.stderr) def set_epoch(self, seed): """Set the current epoch (for consistent shard selection among nodes).""" self.rng = random.Random(seed) def get_shards_for_epoch(self): result = [] for source in self.sources: if source.resample > 0: # sample with replacement l = self.rng.choices(source.urls, k=source.resample) elif source.perepoch > 0: # sample without replacement l = list(source.urls) self.rng.shuffle(l) l = l[: source.perepoch] else: l = list(source.urls) result += l self.rng.shuffle(result) return result def __iter__(self): shards = self.get_shards_for_epoch() for shard in shards: yield dict(url=shard)
An iterable Dataset.
All datasets that represent an iterable of data samples should subclass it. Such form of datasets is particularly useful when data come from a stream.
All subclasses should overwrite :meth:__iter__
, which would return an
iterator of samples in this dataset.
When a subclass is used with :class:~torch.utils.data.DataLoader
, each
item in the dataset will be yielded from the :class:~torch.utils.data.DataLoader
iterator. When :attr:num_workers > 0
, each worker process will have a
different copy of the dataset object, so it is often desired to configure
each copy independently to avoid having duplicate data returned from the
workers. :func:~torch.utils.data.get_worker_info
, when called in a worker
process, returns information about the worker. It can be used in either the
dataset's :meth:__iter__
method or the :class:~torch.utils.data.DataLoader
's
:attr:worker_init_fn
option to modify each copy's behavior.
Example 1: splitting workload across all workers in :meth:__iter__
::
>>> class MyIterableDataset(torch.utils.data.IterableDataset):
... def __init__(self, start, end):
... super(MyIterableDataset).__init__()
... assert end > start, "this example code only works with end >= start"
... self.start = start
... self.end = end
...
... def __iter__(self):
... worker_info = torch.utils.data.get_worker_info()
... if worker_info is None: # single-process data loading, return the full iterator
... iter_start = self.start
... iter_end = self.end
... else: # in a worker process
... # split workload
... per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
... worker_id = worker_info.id
... iter_start = self.start + worker_id * per_worker
... iter_end = min(iter_start + per_worker, self.end)
... return iter(range(iter_start, iter_end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)
>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[3, 4, 5, 6]
>>> # Mult-process loading with two worker processes
>>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[3, 5, 4, 6]
>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=20)))
[3, 4, 5, 6]
Example 2: splitting workload across all workers using :attr:worker_init_fn
::
>>> class MyIterableDataset(torch.utils.data.IterableDataset):
... def __init__(self, start, end):
... super(MyIterableDataset).__init__()
... assert end > start, "this example code only works with end >= start"
... self.start = start
... self.end = end
...
... def __iter__(self):
... return iter(range(self.start, self.end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)
>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[3, 4, 5, 6]
>>>
>>> # Directly doing multi-process loading yields duplicate data
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[3, 3, 4, 4, 5, 5, 6, 6]
>>> # Define a `worker_init_fn` that configures each dataset copy differently
>>> def worker_init_fn(worker_id):
... worker_info = torch.utils.data.get_worker_info()
... dataset = worker_info.dataset # the dataset copy in this worker process
... overall_start = dataset.start
... overall_end = dataset.end
... # configure the dataset to only process the split workload
... per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers)))
... worker_id = worker_info.id
... dataset.start = overall_start + worker_id * per_worker
... dataset.end = min(dataset.start + per_worker, overall_end)
...
>>> # Mult-process loading with the custom `worker_init_fn`
>>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn)))
[3, 5, 4, 6]
>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=20, worker_init_fn=worker_init_fn)))
[3, 4, 5, 6]
View Source
def __init__(self, fname): """Construct a shardlist from multiple sources using a YAML spec.""" self.epoch = -1 self.parse_spec(fname)
Construct a shardlist from multiple sources using a YAML spec.
View Source
def parse_spec(self, fname): self.rng = default_rng # capture default_rng if we fork if isinstance(fname, dict): spec = fname fname = "{dict}" else: with open(fname) as stream: spec = yaml.safe_load(stream) assert set(spec.keys()).issubset(set("prefix datasets buckets".split())), list( spec.keys() ) prefix = expand(spec.get("prefix", "")) self.sources = [] for ds in spec["datasets"]: assert set(ds.keys()).issubset( set("buckets name shards resample choose".split()) ), list(ds.keys()) buckets = ds.get("buckets", spec.get("buckets", [])) if isinstance(buckets, str): buckets = [buckets] buckets = [expand(s) for s in buckets] if buckets == []: buckets = [""] assert ( len(buckets) == 1 ), f"{buckets}: FIXME support for multiple buckets unimplemented" bucket = buckets[0] name = ds.get("name", "@" + bucket) urls = ds["shards"] if isinstance(urls, str): urls = [urls] # urls = [u for url in urls for u in braceexpand.braceexpand(url)] urls = [ prefix + os.path.join(bucket, u) for url in urls for u in braceexpand.braceexpand(expand(url)) ] resample = ds.get("resample", -1) nsample = ds.get("choose", -1) if nsample > len(urls): raise ValueError( f"perepoch {nsample} must be no greater than the number of shards" ) if (nsample > 0) and (resample > 0): raise ValueError("specify only one of perepoch or choose") entry = MSSource(name=name, urls=urls, perepoch=nsample, resample=resample) self.sources.append(entry) print(f"# {name} {len(urls)} {nsample}", file=sys.stderr)
View Source
def set_epoch(self, seed): """Set the current epoch (for consistent shard selection among nodes).""" self.rng = random.Random(seed)
Set the current epoch (for consistent shard selection among nodes).
View Source
def get_shards_for_epoch(self): result = [] for source in self.sources: if source.resample > 0: # sample with replacement l = self.rng.choices(source.urls, k=source.resample) elif source.perepoch > 0: # sample without replacement l = list(source.urls) self.rng.shuffle(l) l = l[: source.perepoch] else: l = list(source.urls) result += l self.rng.shuffle(result) return result
Inherited Members
- torch.utils.data.dataset.IterableDataset
- functions
- reduce_ex_hook
- register_function
- register_datapipe_as_function
- set_reduce_ex_hook
- type
View Source
class ResampledShards(IterableDataset): """An iterable dataset yielding a list of urls.""" def __init__( self, urls, nshards=sys.maxsize, worker_seed=None, deterministic=False, ): """Sample shards from the shard list with replacement. :param urls: a list of URLs as a Python list or brace notation string """ super().__init__() urls = expand_urls(urls) self.urls = urls assert isinstance(self.urls[0], str) self.nshards = nshards self.rng = random.Random() self.worker_seed = ( utils.pytorch_worker_seed if worker_seed is None else worker_seed ) self.deterministic = deterministic self.epoch = -1 def __iter__(self): """Return an iterator over the shards.""" self.epoch += 1 if self.deterministic: seed = (self.worker_seed(), self.epoch) else: seed = (self.worker_seed(), self.epoch, os.getpid(), time.time()) self.rng.seed(seed) for _ in range(self.nshards): yield dict(url=self.rng.choice(self.urls))
An iterable dataset yielding a list of urls.
View Source
def __init__( self, urls, nshards=sys.maxsize, worker_seed=None, deterministic=False, ): """Sample shards from the shard list with replacement. :param urls: a list of URLs as a Python list or brace notation string """ super().__init__() urls = expand_urls(urls) self.urls = urls assert isinstance(self.urls[0], str) self.nshards = nshards self.rng = random.Random() self.worker_seed = ( utils.pytorch_worker_seed if worker_seed is None else worker_seed ) self.deterministic = deterministic self.epoch = -1
Sample shards from the shard list with replacement.
:param urls: a list of URLs as a Python list or brace notation string
Inherited Members
- torch.utils.data.dataset.IterableDataset
- functions
- reduce_ex_hook
- register_function
- register_datapipe_as_function
- set_reduce_ex_hook
- type