Dataset mixture at the shard level

README.md

@@ -197,8 +197,10 @@ Using `--dataset-resampled` is recommended for these cases.
 By default, on expectation the amount of times the model will see a sample from each source is proportional to the size of the source.
 For instance, when training on one data source with size 400M and one with size 10M, samples from the first source are 40x more likely to be seen in expectation.
 
-We also support different weighting of the data sources, by using the `--train-data-weights` flag.
-For instance, using `--train-data-weights=1::1` in the above scenario will make it so that samples from each data source are equally likely to be seen in expectation. 
+We also support different weighting of the data sources, by using the `--train_data_upsampling_factors` flag.
+For instance, using `--train_data_upsampling_factors=1::1` in the above scenario is equivalent to not using the flag, and `--train_data_upsampling_factors=1::2` is equivalent to upsampling the second data source twice.
+If you want to sample from data sources with the same frequency, the upsampling factors should be inversely proportional to the sizes of the data sources.
+For instance, if dataset `A` has 1000 samples and dataset `B` has 100 samples, you can use `--train_data_upsampling_factors=0.001::0.01` (or analogously, ``--train_data_upsampling_factors=1::10`).
 
 #### Single-Node
 

src/training/data.py

@@ -6,6 +6,7 @@
 import random
 import sys
 import time
+import braceexpand
 from dataclasses import dataclass
 from multiprocessing import Value
 
@@ -71,8 +72,29 @@ def set_epoch(self, epoch):
             self.sampler.set_epoch(epoch)
 
 
+def expand_urls(urls, weights=None):
+    if weights is None:
+        expanded_urls = wds.shardlists.expand_urls(urls)
+        return expanded_urls, None
+    if isinstance(urls, str):
+        urllist = urls.split("::")
+        weights = weights.split('::')
+        assert len(weights) == len(urllist), f"Expected the number of data components ({len(urllist)}) and weights({len(weights)}) to match."
+        weights = [float(weight) for weight in weights]
+        all_urls, all_weights = [], []
+        for url, weight in zip(urllist, weights):
+            expanded_url = list(braceexpand.braceexpand(url))
+            expanded_weights = [weight for _ in expanded_url]
+            all_urls.extend(expanded_url)
+            all_weights.extend(expanded_weights)
+        return all_urls, all_weights
+    else:
+        all_urls = list(urls)
+        return all_urls, weights
+
+
 def get_dataset_size(shards):
-    shards_list = wds.shardlists.expand_urls(shards)
+    shards_list, _ = expand_urls(shards)
     dir_path = os.path.dirname(shards_list[0])
     sizes_filename = os.path.join(dir_path, 'sizes.json')
     len_filename = os.path.join(dir_path, '__len__')
@@ -255,6 +277,7 @@ class ResampledShards2(IterableDataset):
     def __init__(
         self,
         urls,
+        weights=None,
         nshards=sys.maxsize,
         worker_seed=None,
         deterministic=False,
@@ -265,8 +288,11 @@ def __init__(
         :param urls: a list of URLs as a Python list or brace notation string
         """
         super().__init__()
-        urls = wds.shardlists.expand_urls(urls)
+        urls, weights = expand_urls(urls, weights)
         self.urls = urls
+        self.weights = weights
+        if self.weights is not None:
+            assert len(self.urls) == len(self.weights), f"Number of urls {len(self.urls)} and weights {len(self.weights)} should match."
         assert isinstance(self.urls[0], str)
         self.nshards = nshards
         self.rng = random.Random()
@@ -292,28 +318,10 @@ def __iter__(self):
                 seed = self.worker_seed() + epoch
             self.rng.seed(seed)
         for _ in range(self.nshards):
-            yield dict(url=self.rng.choice(self.urls))
-
-
-class WeightedSampler(IterableDataset):
-    def __init__(self, datasets, weights=None):
-        super().__init__()
-        self.datasets = datasets
-        self.rng = random.Random()
-        if weights is None:
-            weights = [1 for _ in self.datasets]
-        self.weights = weights
-
-    def __iter__(self):
-        sources = [iter(ds) for ds in self.datasets]
-        idxs = list(range(len(sources)))
-        for _ in range(sys.maxsize):
-            idx = self.rng.choices(idxs, weights=self.weights, k=1)[0]
-            try:
-                sample = next(sources[idx])
-                yield sample
-            except StopIteration:
-                return
+            if self.weights is None:
+                yield dict(url=self.rng.choice(self.urls))
+            else:
+                yield dict(url=self.rng.choices(self.urls, weights=self.weights, k=1)[0])
 
 
 def get_wds_dataset(args, preprocess_img, is_train, epoch=0, floor=False, tokenizer=None):
@@ -334,62 +342,47 @@ def get_wds_dataset(args, preprocess_img, is_train, epoch=0, floor=False, tokeni
 
     shared_epoch = SharedEpoch(epoch=epoch)  # create a shared epoch store to sync epoch to dataloader worker proc
 
-    def build_pipeline(input_shards):
-        if resampled:
-            pipeline = [ResampledShards2(input_shards, deterministic=True, epoch=shared_epoch)]
-        else:
-            pipeline = [wds.SimpleShardList(input_shards)]
+    if resampled:
+        pipeline = [ResampledShards2(input_shards, weights=args.train_data_upsampling_factors, deterministic=True, epoch=shared_epoch)]
+    else:
+        assert args.train_data_upsampling_factors is None, "--train_data_upsampling_factors is only supported when sampling with replacement (together with --dataset-resampled)."
+        pipeline = [wds.SimpleShardList(input_shards)]
 
-        # at this point we have an iterator over all the shards
-        if is_train:
-            if not resampled:
-                pipeline.extend([
-                    detshuffle2(
-                        bufsize=_SHARD_SHUFFLE_SIZE,
-                        initial=_SHARD_SHUFFLE_INITIAL,
-                        seed=args.seed,
-                        epoch=shared_epoch,
-                    ),
-                    wds.split_by_node,
-                    wds.split_by_worker,
-                ])
+    # at this point we have an iterator over all the shards
+    if is_train:
+        if not resampled:
             pipeline.extend([
-                # at this point, we have an iterator over the shards assigned to each worker at each node
-                tarfile_to_samples_nothrow,  # wds.tarfile_to_samples(handler=log_and_continue),
-                wds.shuffle(
-                    bufsize=_SAMPLE_SHUFFLE_SIZE,
-                    initial=_SAMPLE_SHUFFLE_INITIAL,
+                detshuffle2(
+                    bufsize=_SHARD_SHUFFLE_SIZE,
+                    initial=_SHARD_SHUFFLE_INITIAL,
+                    seed=args.seed,
+                    epoch=shared_epoch,
                 ),
-            ])
-        else:
-            pipeline.extend([
+                wds.split_by_node,
                 wds.split_by_worker,
-                # at this point, we have an iterator over the shards assigned to each worker
-                wds.tarfile_to_samples(handler=log_and_continue),
             ])
         pipeline.extend([
-            wds.select(filter_no_caption_or_no_image),
-            wds.decode("pilrgb", handler=log_and_continue),
-            wds.rename(image="jpg;png;jpeg;webp", text="txt"),
-            wds.map_dict(image=preprocess_img, text=lambda text: tokenizer(text)[0]),
-            wds.to_tuple("image", "text"),
+            # at this point, we have an iterator over the shards assigned to each worker at each node
+            tarfile_to_samples_nothrow,  # wds.tarfile_to_samples(handler=log_and_continue),
+            wds.shuffle(
+                bufsize=_SAMPLE_SHUFFLE_SIZE,
+                initial=_SAMPLE_SHUFFLE_INITIAL,
+            ),
         ])
-
-        return pipeline
-
-    if args.train_data_weights is not None:
-        weights = args.train_data_weights
-        weights = [float(w) for w in weights.split('::')]
-        input_shard_list = input_shards.split('::')
-        assert len(input_shard_list) == len(weights), f"Expected the number of data components ({len(input_shard_list)}) and weights ({len(weights)}) to match."
-        datasets = [wds.DataPipeline(*build_pipeline(shards)) for shards in input_shard_list]
-        pipeline = [WeightedSampler(datasets, weights=weights)]
     else:
-        pipeline = build_pipeline(input_shards)
-
-    pipeline.append(
+        pipeline.extend([
+            wds.split_by_worker,
+            # at this point, we have an iterator over the shards assigned to each worker
+            wds.tarfile_to_samples(handler=log_and_continue),
+        ])
+    pipeline.extend([
+        wds.select(filter_no_caption_or_no_image),
+        wds.decode("pilrgb", handler=log_and_continue),
+        wds.rename(image="jpg;png;jpeg;webp", text="txt"),
+        wds.map_dict(image=preprocess_img, text=lambda text: tokenizer(text)[0]),
+        wds.to_tuple("image", "text"),
         wds.batched(args.batch_size, partial=not is_train)
-    )
+    ])
 
     dataset = wds.DataPipeline(*pipeline)
 

src/training/params.py

@@ -32,11 +32,11 @@ def parse_args(args):
         help="Path to file(s) with training data. When using webdataset, multiple datasources can be combined using the `::` separator.",
     )
     parser.add_argument(
-        "--train-data-weights",
+        "--train-data-upsampling-factors",
         type=str,
         default=None,
         help=(
-            "When using multiple data sources with webdataset and sampling with replacement, which weight to use for sampling the different data sources. "
+            "When using multiple data sources with webdataset and sampling with replacement, this can be used to upsample specific data sources. "
             "Similar to --train-data, this should be a string with as many numbers as there are data sources, separated by `::` (e.g. 1::2::0.5) "
             "By default, datapoints are sampled uniformly regardless of the dataset sizes."
         )

tests/test_wds.py

@@ -0,0 +1,149 @@
+import os
+import pytest
+import util_test
+import collections
+import tarfile
+import io
+from PIL import Image
+
+from training.data import get_wds_dataset
+from training.params import parse_args
+from training.main import random_seed
+
+TRAIN_NUM_SAMPLES = 10_000
+RTOL = 0.2
+
+# NOTE: we use two test tar files, which are created on the fly and saved to data/input.
+# 000.tar has 10 samples, and the captions are 000_0, 000_1, ..., 000_9
+# 001.tar has 5 samples, and the captions are 001_0, 001_1, ..., 001_4
+def build_inputs(test_name):
+    base_input_dir, _ = util_test.get_data_dirs()
+    input_dir = os.path.join(base_input_dir, test_name)
+    os.makedirs(input_dir, exist_ok=True)
+
+    def save_tar(idx, num_samples):
+        filename = os.path.join(input_dir, f'test_data_{idx:03d}.tar')
+        tar = tarfile.open(filename, 'w')
+
+        for sample_idx in range(num_samples):
+            # Image
+            image = Image.new('RGB', (32, 32))
+            info = tarfile.TarInfo(f'{sample_idx}.png')
+            bio = io.BytesIO()
+            image.save(bio, format='png')
+            size = bio.tell()
+            bio.seek(0)
+            info.size = size
+            tar.addfile(info, bio)
+
+            # Caption
+            info = tarfile.TarInfo(f'{sample_idx}.txt')
+            bio = io.BytesIO()
+            bio.write(f'{idx:03d}_{sample_idx}'.encode('utf-8'))
+            size = bio.tell()
+            bio.seek(0)
+            info.size = size
+            tar.addfile(info, bio)
+
+        tar.close()          
+
+    save_tar(0, 10)
+    save_tar(1, 5)
+
+    return input_dir
+
+
+def build_params(input_shards, seed=0):
+    args = parse_args([])
+    args.train_data = input_shards
+    args.train_num_samples = TRAIN_NUM_SAMPLES
+    args.dataset_resampled = True
+    args.seed = seed
+    args.workers = 1
+    args.world_size = 1
+    args.batch_size = 1
+    random_seed(seed)
+
+    preprocess_img = lambda x: x
+    tokenizer = lambda x: [x.strip()]
+
+    return args, preprocess_img, tokenizer
+
+
+def get_dataloader(input_shards):
+    args, preprocess_img, tokenizer = build_params(input_shards)
+    dataset = get_wds_dataset(args, preprocess_img, is_train=True, tokenizer=tokenizer)
+    dataloader = dataset.dataloader
+    return dataloader
+
+
+def test_single_source():
+    """Test webdataset with a single tar file."""
+    input_dir = build_inputs('single_source')    
+    input_shards = os.path.join(input_dir, 'test_data_000.tar')
+    dataloader = get_dataloader(input_shards)
+
+    counts = collections.defaultdict(int)
+    for sample in dataloader:
+        txts = sample[1]
+        for txt in txts:
+            counts[txt] += 1
+
+    for key, count in counts.items():
+        assert count == pytest.approx(TRAIN_NUM_SAMPLES / 10, RTOL)
+
+
+def test_two_sources():
+    """Test webdataset with a single two tar files."""
+    input_dir = build_inputs('two_sources')
+    input_shards = os.path.join(input_dir, 'test_data_{000..001}.tar')
+    dataloader = get_dataloader(input_shards)
+
+    counts = collections.defaultdict(int)
+    for sample in dataloader:
+        txts = sample[1]
+        for txt in txts:
+            counts[txt] += 1
+
+    for key, count in counts.items():
+        assert count == pytest.approx(TRAIN_NUM_SAMPLES / 15, RTOL), f'{key}, {count}'
+
+
+def test_two_sources_same_weights():
+    """Test webdataset with a two tar files, using --train-data-weights=1::1."""
+    input_dir = build_inputs('two_sources_same_weights')
+    input_shards = f"{os.path.join(input_dir, 'test_data_000.tar')}::{os.path.join(input_dir, 'test_data_001.tar')}"
+    args, preprocess_img, tokenizer = build_params(input_shards)
+    args.train_data_upsampling_factors = '1::1'
+    dataset = get_wds_dataset(args, preprocess_img, is_train=True, tokenizer=tokenizer)
+    dataloader = dataset.dataloader
+
+    counts = collections.defaultdict(int)
+    for sample in dataloader:
+        txts = sample[1]
+        for txt in txts:
+            counts[txt] += 1
+
+    for key, count in counts.items():
+        assert count == pytest.approx(TRAIN_NUM_SAMPLES / 15, RTOL), f'{key}, {count}'
+
+def test_two_sources_with_upsampling():
+    """Test webdataset with a two tar files with upsampling."""
+    input_dir = build_inputs('two_sources_with_upsampling')
+    input_shards = f"{os.path.join(input_dir, 'test_data_000.tar')}::{os.path.join(input_dir, 'test_data_001.tar')}"
+    args, preprocess_img, tokenizer = build_params(input_shards)
+    args.train_data_upsampling_factors = '1::2'
+    dataset = get_wds_dataset(args, preprocess_img, is_train=True, tokenizer=tokenizer)
+    dataloader = dataset.dataloader
+
+    counts = collections.defaultdict(int)
+    for sample in dataloader:
+        txts = sample[1]
+        for txt in txts:
+            counts[txt] += 1
+
+    for key, count in counts.items():
+        if key.startswith('000'):
+            assert count == pytest.approx(TRAIN_NUM_SAMPLES / 20, RTOL), f'{key}, {count}'
+        else:
+            assert count == pytest.approx(TRAIN_NUM_SAMPLES / 10, RTOL), f'{key}, {count}'