Fix up a few issues with variable image/batch size scheduler

Author: rwightman

tests/test_naflex_dataset.py

@@ -1,11 +1,13 @@
+import warnings
+from types import SimpleNamespace
+
 import torch
 from torch.utils.data import DataLoader, Dataset
 
 from timm.data import NaFlexMapDatasetWrapper
 
 
 class _TensorImageDataset(Dataset):
-
     def __init__(self, length=32):
         self.length = length
 
@@ -68,19 +70,29 @@ def test_naflex_persistent_workers_read_shared_epoch():
         dataset.set_epoch(1)
         epoch_1_indices = _loader_indices(loader)
 
-        expected_epoch_0 = [
-            index
-            for _, _, indices in dataset._prepare_epoch_batches(0)
-            for index in indices
-        ]
-        expected_epoch_1 = [
-            index
-            for _, _, indices in dataset._prepare_epoch_batches(1)
-            for index in indices
-        ]
+        expected_epoch_0 = [index for _, _, indices in dataset._prepare_epoch_batches(0) for index in indices]
+        expected_epoch_1 = [index for _, _, indices in dataset._prepare_epoch_batches(1) for index in indices]
         assert epoch_0_indices == expected_epoch_0
         assert epoch_1_indices == expected_epoch_1
         assert epoch_1_indices != epoch_0_indices
     finally:
         if loader._iterator is not None:
             loader._iterator._shutdown_workers()
+
+
+def test_naflex_epoch_prep_warnings_only_emitted_by_worker_zero(monkeypatch):
+    dataset = _create_naflex_dataset()
+
+    def prepare(_epoch):
+        warnings.warn('epoch prep mismatch')
+        return []
+
+    dataset._prepare_epoch_batches = prepare
+
+    for worker_id, expected_warnings in ((0, 1), (1, 0)):
+        worker_info = SimpleNamespace(id=worker_id, num_workers=2)
+        monkeypatch.setattr(torch.utils.data, 'get_worker_info', lambda: worker_info)
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter('always')
+            list(dataset)
+        assert len(caught) == expected_warnings

tests/test_scheduled_sampler.py

@@ -9,7 +9,6 @@
 
 
 class _ImageDataset(Dataset):
-
     def __init__(self, length=64, image_size=48):
         self.length = length
         self.image_size = image_size
@@ -55,7 +54,6 @@ def test_scheduled_batch_sampler_stable_epoch_length_and_composition():
 
 def test_sample_budget_schedule_is_created_once_and_cached():
     class TrackingScheduledBatchSampler(ScheduledBatchSampler):
-
         def __init__(self, *args, **kwargs):
             self.schedule_create_count = 0
             super().__init__(*args, **kwargs)
@@ -75,6 +73,60 @@ def _create_sample_budget_schedule(self):
     assert batch_sampler.schedule_create_count == 1
 
 
+def test_zero_weight_choices_are_excluded_from_progressive_random_mix_and_sample_budget():
+    progressive_sampler = ScheduledBatchSampler(
+        SequentialSampler(range(64)),
+        batch_sizes=(16, 8, 4),
+        choice_weights=(1, 0, 0),
+        choice_schedule='progressive',
+        schedule_epochs=3,
+        schedule_random_mix=0.1,
+    )
+
+    for epoch in range(3):
+        assert progressive_sampler.choice_weights_for_epoch(epoch)[1:].tolist() == [0, 0]
+
+    snap_sampler = ScheduledBatchSampler(
+        SequentialSampler(range(64)),
+        batch_sizes=(16, 12, 8, 4),
+        choice_weights=(1, 0, 0, 1),
+        choice_schedule='progressive',
+        schedule_epochs=4,
+        schedule_spread=0,
+        schedule_random_mix=0,
+    )
+    assert snap_sampler.choice_weights_for_epoch(1).tolist() == [1, 0, 0, 0]
+    assert snap_sampler.choice_weights_for_epoch(2).tolist() == [0, 0, 0, 1]
+
+    sample_budget_sampler = ScheduledBatchSampler(
+        SequentialSampler(range(12)),
+        batch_sizes=(8, 4),
+        choice_weights=(1, 0),
+    )
+    batches = list(sample_budget_sampler)
+    assert _batch_signature(batches) == [(0, 8)]
+
+
+def test_progressive_schedule_rejects_when_no_full_batch_fits():
+    with pytest.raises(ValueError, match='No full scheduled batch'):
+        ScheduledBatchSampler(
+            SequentialSampler(range(3)),
+            batch_sizes=(8, 4),
+            choice_schedule='progressive',
+            schedule_epochs=3,
+        )
+
+
+def test_sample_budget_schedule_rejects_when_no_full_batch_fits():
+    with pytest.raises(ValueError, match='No full scheduled batch'):
+        ScheduledBatchSampler(SequentialSampler(range(3)), batch_sizes=(8, 4))
+
+
+def test_scheduled_batch_sampler_rejects_empty_sampler():
+    with pytest.raises(ValueError, match='non-empty sampler'):
+        ScheduledBatchSampler(SequentialSampler(range(0)), batch_sizes=(8, 4))
+
+
 def test_scheduled_batch_sampler_distributed_shapes_match():
     dataset = list(range(96))
     rank_0_sampler = DistributedSampler(dataset, num_replicas=2, rank=0, shuffle=True, seed=11)
@@ -123,13 +175,13 @@ def test_progressive_schedule_moves_choices_and_preserves_constant_batch_budget(
     assert [index for batch in epoch_0 for index, _ in batch] == list(range(1000))
     assert [index for batch in epoch_4 for index, _ in batch] == list(range(1000))
     assert sum(choice for choice, _ in _batch_signature(epoch_0)) < sum(
-        choice for choice, _ in _batch_signature(epoch_4)
+        choice
+        for choice, _ in _batch_signature(epoch_4)
     )
 
 
 def test_progressive_schedule_is_created_when_iteration_starts():
     class TrackingScheduledBatchSampler(ScheduledBatchSampler):
-
         def __init__(self, *args, **kwargs):
             self.created_epochs = []
             super().__init__(*args, **kwargs)
@@ -168,19 +220,21 @@ def test_progressive_schedule_infers_policy_average_batch_budget_and_cycles_indi
         schedule_random_mix=0.1,
     )
     batch_sizes = torch.tensor(batch_sampler.batch_sizes, dtype=torch.float64)
-    expected_average = torch.stack([
-        torch.dot(batch_sampler.choice_weights_for_epoch(epoch), batch_sizes)
-        for epoch in range(5)
-    ]).mean().item()
+    expected_average = (
+        torch
+        .stack([torch.dot(batch_sampler.choice_weights_for_epoch(epoch), batch_sizes) for epoch in range(5)])
+        .mean()
+        .item()
+    )
 
     assert batch_sampler.average_batch_size == pytest.approx(expected_average)
     assert len(batch_sampler) == int(len(sampler) / expected_average)
 
     epoch_0 = list(batch_sampler)
     epoch_0_indices = [index for batch in epoch_0 for index, _ in batch]
     assert len(epoch_0_indices) > len(sampler)
-    assert epoch_0_indices[:len(sampler)] == list(range(len(sampler)))
-    assert epoch_0_indices[len(sampler):] == list(range(len(epoch_0_indices) - len(sampler)))
+    assert epoch_0_indices[: len(sampler)] == list(range(len(sampler)))
+    assert epoch_0_indices[len(sampler) :] == list(range(len(epoch_0_indices) - len(sampler)))
 
     batch_sampler.set_epoch(4)
     epoch_4 = list(batch_sampler)
@@ -190,10 +244,7 @@ def test_progressive_schedule_infers_policy_average_batch_budget_and_cycles_indi
 
 def test_progressive_schedule_distributed_shapes_match_and_advance():
     dataset = list(range(96))
-    samplers = [
-        DistributedSampler(dataset, num_replicas=2, rank=rank, shuffle=True, seed=11)
-        for rank in range(2)
-    ]
+    samplers = [DistributedSampler(dataset, num_replicas=2, rank=rank, shuffle=True, seed=11) for rank in range(2)]
     batch_samplers = [
         ScheduledBatchSampler(
             sampler,
@@ -220,16 +271,12 @@ def test_progressive_schedule_distributed_shapes_match_and_advance():
     for rank_batches in (*epoch_0, *epoch_2):
         assert len(rank_batches) == 8
     for rank_epochs in (epoch_0, epoch_2):
-        rank_indices = [
-            {index for batch in rank_batches for index, _ in batch}
-            for rank_batches in rank_epochs
-        ]
+        rank_indices = [{index for batch in rank_batches for index, _ in batch} for rank_batches in rank_epochs]
         assert rank_indices[0].isdisjoint(rank_indices[1])
 
 
 def test_scheduled_transform_dataset_preserves_sample_fields():
     class ExtraFieldDataset(Dataset):
-
         def __getitem__(self, index):
             return index, index + 1, f'sample-{index}'
 
@@ -349,6 +396,30 @@ def test_create_loader_rejects_scheduled_resolutions_with_multi_epochs_loader():
         )
 
 
+@pytest.mark.parametrize(
+    'scheduled_option',
+    (
+        {'batch_choice_seed': 7},
+        {'batch_schedule_epochs': 3},
+        {'batch_schedule_spread': 0.5},
+        {'batch_schedule_random_mix': 0.2},
+    ),
+)
+def test_create_loader_rejects_scheduled_options_without_input_size_choices(scheduled_option):
+    with pytest.raises(ValueError, match='input_size_choices'):
+        create_loader(
+            _ImageDataset(length=32),
+            input_size=(3, 32, 32),
+            batch_size=4,
+            is_training=True,
+            no_aug=True,
+            use_prefetcher=False,
+            num_workers=0,
+            persistent_workers=False,
+            **scheduled_option,
+        )
+
+
 def test_create_loader_standard_batching_path_is_unchanged():
     dataset = _ImageDataset()
     loader = create_loader(

timm/data/loader.py

@@ -302,7 +302,8 @@ def create_loader(
         batch_choice_schedule: Choice schedule mode, either ``constant`` or ``progressive``.
         batch_schedule_epochs: Number of epochs over which a progressive schedule moves through its choices.
         batch_schedule_spread: Standard deviation of the progressive choice window, in choice-index units.
-        batch_schedule_random_mix: Fraction of uniform random exploration in a progressive schedule.
+        batch_schedule_random_mix: Fraction of uniform random exploration over nonzero-weight choices in a
+            progressive schedule.
         num_batches: Fixed number of loader batches per epoch. Inferred from the schedule-average batch size
             for progressive schedules when not specified.
 
@@ -387,7 +388,11 @@ def create_loader(
         if (
                 batch_size_choices is not None
                 or batch_choice_weights is not None
+                or batch_choice_seed != 0
                 or batch_choice_schedule != 'constant'
+                or batch_schedule_epochs is not None
+                or batch_schedule_spread != 0.65
+                or batch_schedule_random_mix != 0.1
                 or num_batches is not None
         ):
             raise ValueError('input_size_choices must be specified when using scheduled batch options.')

timm/data/naflex_dataset.py

@@ -512,7 +512,13 @@ def __iter__(self) -> Iterator[Tuple[Dict[str, torch.Tensor], torch.Tensor]]:
 
         # Snapshot the shared epoch once, then build process-local derived state.
         epoch = self.shared_epoch.value
-        batches_for_worker = self._prepare_epoch_batches(epoch)[worker_id::num_workers]
+        if worker_id == 0:
+            epoch_batches = self._prepare_epoch_batches(epoch)
+        else:
+            with warnings.catch_warnings():
+                warnings.simplefilter('ignore')
+                epoch_batches = self._prepare_epoch_batches(epoch)
+        batches_for_worker = epoch_batches[worker_id::num_workers]
         for seq_len, patch_idx, indices in batches_for_worker:
             if not indices: # Skip if a batch ended up with no indices (shouldn't happen often)
                  continue