diff --git a/tests/test_naflex_dataset.py b/tests/test_naflex_dataset.py new file mode 100644 index 0000000000..ce23b697d1 --- /dev/null +++ b/tests/test_naflex_dataset.py @@ -0,0 +1,98 @@ +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 + + def __getitem__(self, index): + return torch.full((3, 2, 2), index, dtype=torch.float32), index + + def __len__(self): + return self.length + + +def _create_naflex_dataset(epoch=0): + return NaFlexMapDatasetWrapper( + _TensorImageDataset(), + patch_size=1, + seq_lens=(1,), + max_tokens_per_batch=4, + seed=17, + epoch=epoch, + batch_divisor=1, + ) + + +def _loader_indices(loader): + return [index for _, targets in loader for index in targets.tolist()] + + +def test_naflex_epoch_batches_are_prepared_when_iteration_starts(): + dataset = _create_naflex_dataset() + prepare_epoch_batches = dataset._prepare_epoch_batches + prepared_epochs = [] + + def prepare(epoch): + prepared_epochs.append(epoch) + return prepare_epoch_batches(epoch) + + dataset._prepare_epoch_batches = prepare + + assert prepared_epochs == [] + dataset.set_epoch(3) + assert prepared_epochs == [] + assert dataset.shared_epoch.value == 3 + + iterator = iter(dataset) + assert prepared_epochs == [] + next(iterator) + assert prepared_epochs == [3] + + +def test_naflex_persistent_workers_read_shared_epoch(): + dataset = _create_naflex_dataset() + loader = DataLoader( + dataset, + batch_size=None, + num_workers=2, + persistent_workers=True, + ) + + try: + epoch_0_indices = _loader_indices(loader) + 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] + 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 diff --git a/tests/test_scheduled_sampler.py b/tests/test_scheduled_sampler.py new file mode 100644 index 0000000000..46dccacc9c --- /dev/null +++ b/tests/test_scheduled_sampler.py @@ -0,0 +1,480 @@ +from collections import Counter + +import pytest +import torch +from PIL import Image +from torch.utils.data import Dataset, DistributedSampler, SequentialSampler + +from timm.data import FastCollateMixup, ScheduledBatchSampler, ScheduledTransformDataset, create_loader + + +class _ImageDataset(Dataset): + def __init__(self, length=64, image_size=48): + self.length = length + self.image_size = image_size + self.transform = None + + def __getitem__(self, index): + image = Image.new('RGB', (self.image_size, self.image_size), color=index % 256) + if self.transform is not None: + image = self.transform(image) + return image, index + + def __len__(self): + return self.length + + +def _batch_signature(batches): + return [(batch[0][1], len(batch)) for batch in batches] + + +def test_scheduled_batch_sampler_stable_epoch_length_and_composition(): + sampler = SequentialSampler(range(101)) + batch_sampler = ScheduledBatchSampler( + sampler, + batch_sizes=(16, 8, 4), + choice_weights=(0.2, 0.3, 0.5), + seed=17, + ) + + epoch_0 = list(batch_sampler) + batch_sampler.set_epoch(1) + epoch_1 = list(batch_sampler) + + assert len(epoch_0) == len(epoch_1) == len(batch_sampler) + assert Counter(_batch_signature(epoch_0)) == Counter(_batch_signature(epoch_1)) + assert _batch_signature(epoch_0) != _batch_signature(epoch_1) + assert sum(map(len, epoch_0)) == sum(batch_size for _, batch_size in batch_sampler.schedule) + assert 0 <= len(sampler) - sum(map(len, epoch_0)) < min(batch_sampler.batch_sizes) + assert [sample_index for batch in epoch_0 for sample_index, _ in batch] == list( + range(sum(map(len, epoch_0))) + ) + assert all(len({choice for _, choice in batch}) == 1 for batch in epoch_0) + + +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) + + def _create_sample_budget_schedule(self): + self.schedule_create_count += 1 + return super()._create_sample_budget_schedule() + + batch_sampler = TrackingScheduledBatchSampler( + SequentialSampler(range(101)), + batch_sizes=(16, 8, 4), + seed=17, + ) + + assert batch_sampler.schedule_create_count == 1 + assert len(batch_sampler) == len(batch_sampler) == len(list(batch_sampler)) + assert batch_sampler.schedule_create_count == 1 + + +def test_constant_schedule_ignores_progressive_only_options(): + batch_sampler = ScheduledBatchSampler( + SequentialSampler(range(32)), + batch_sizes=(8, 4), + choice_schedule='constant', + schedule_epochs=0, + schedule_spread=-1, + schedule_random_mix=2, + ) + + assert len(batch_sampler) > 0 + + +def test_progressive_schedule_requires_multiple_choices(): + with pytest.raises(ValueError, match='at least two'): + ScheduledBatchSampler( + SequentialSampler(range(32)), + batch_sizes=(8,), + choice_schedule='progressive', + schedule_epochs=3, + ) + + +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) + rank_1_sampler = DistributedSampler(dataset, num_replicas=2, rank=1, shuffle=True, seed=11) + rank_0_batches = ScheduledBatchSampler(rank_0_sampler, (12, 6), seed=29) + rank_1_batches = ScheduledBatchSampler(rank_1_sampler, (12, 6), seed=29) + + rank_0_batches.set_epoch(3) + rank_1_batches.set_epoch(3) + rank_0_epoch = list(rank_0_batches) + rank_1_epoch = list(rank_1_batches) + + assert _batch_signature(rank_0_epoch) == _batch_signature(rank_1_epoch) + rank_0_indices = {index for batch in rank_0_epoch for index, _ in batch} + rank_1_indices = {index for batch in rank_1_epoch for index, _ in batch} + assert rank_0_indices.isdisjoint(rank_1_indices) + + +def test_progressive_schedule_moves_choices_and_preserves_constant_batch_budget(): + sampler = SequentialSampler(range(1000)) + batch_sampler = ScheduledBatchSampler( + sampler, + batch_sizes=(10, 10, 10), + seed=23, + choice_schedule='progressive', + schedule_epochs=5, + schedule_spread=0.5, + schedule_random_mix=0.1, + ) + + assert batch_sampler.average_batch_size == 10 + assert len(batch_sampler) == 100 + start_weights = batch_sampler.choice_weights_for_epoch(0) + middle_weights = batch_sampler.choice_weights_for_epoch(2) + end_weights = batch_sampler.choice_weights_for_epoch(4) + assert start_weights.argmax().item() == 0 + assert middle_weights.argmax().item() == 1 + assert end_weights.argmax().item() == 2 + assert (start_weights >= 0.1 / 3).all() + + epoch_0 = list(batch_sampler) + batch_sampler.set_epoch(4) + epoch_4 = list(batch_sampler) + assert len(epoch_0) == len(epoch_4) == 100 + assert sum(map(len, epoch_0)) == sum(map(len, epoch_4)) == 1000 + 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) + ) + + +def test_progressive_schedule_is_created_when_iteration_starts(): + class TrackingScheduledBatchSampler(ScheduledBatchSampler): + def __init__(self, *args, **kwargs): + self.created_epochs = [] + super().__init__(*args, **kwargs) + + def _create_fixed_batch_schedule(self, epoch): + self.created_epochs.append(epoch) + return super()._create_fixed_batch_schedule(epoch) + + batch_sampler = TrackingScheduledBatchSampler( + SequentialSampler(range(32)), + batch_sizes=(8, 4), + num_batches=4, + choice_schedule='progressive', + schedule_epochs=3, + ) + + assert batch_sampler.created_epochs == [] + batch_sampler.set_epoch(2) + assert batch_sampler.created_epochs == [] + + iterator = iter(batch_sampler) + assert batch_sampler.created_epochs == [] + next(iterator) + assert batch_sampler.created_epochs == [2] + + +def test_progressive_schedule_infers_policy_average_batch_budget_and_cycles_indices(): + sampler = SequentialSampler(range(1000)) + batch_sampler = ScheduledBatchSampler( + sampler, + batch_sizes=(80, 40, 20), + seed=3, + choice_schedule='progressive', + schedule_epochs=5, + schedule_spread=0.5, + 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() + ) + + 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))) + + batch_sampler.set_epoch(4) + epoch_4 = list(batch_sampler) + assert len(epoch_4) == len(epoch_0) + assert sum(map(len, epoch_4)) < len(sampler) + + +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)] + batch_samplers = [ + ScheduledBatchSampler( + sampler, + batch_sizes=(12, 6, 3), + seed=29, + num_batches=8, + choice_schedule='progressive', + schedule_epochs=3, + schedule_spread=0, + schedule_random_mix=0, + ) + for sampler in samplers + ] + + epoch_0 = [list(batch_sampler) for batch_sampler in batch_samplers] + assert _batch_signature(epoch_0[0]) == _batch_signature(epoch_0[1]) + assert {choice for choice, _ in _batch_signature(epoch_0[0])} == {0} + for batch_sampler in batch_samplers: + batch_sampler.set_epoch(2) + epoch_2 = [list(batch_sampler) for batch_sampler in batch_samplers] + assert _batch_signature(epoch_2[0]) == _batch_signature(epoch_2[1]) + assert {choice for choice, _ in _batch_signature(epoch_2[0])} == {2} + + 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] + 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}' + + def __len__(self): + return 4 + + dataset = ScheduledTransformDataset(ExtraFieldDataset(), (lambda value: value * 2, lambda value: value * 3)) + + assert dataset[(2, 0)] == (4, 3, 'sample-2') + assert dataset[(2, 1)] == (6, 3, 'sample-2') + + +def test_create_loader_scheduled_resolutions_and_batch_sizes(): + loader = create_loader( + _ImageDataset(), + input_size=(3, 32, 32), + batch_size=4, + input_size_choices=(16, 32), + batch_size_choices=(8, 4), + batch_choice_weights=(0.5, 0.5), + batch_choice_seed=7, + is_training=True, + no_aug=True, + use_prefetcher=False, + num_workers=0, + persistent_workers=False, + ) + + assert isinstance(loader.dataset, ScheduledTransformDataset) + assert isinstance(loader.batch_sampler, ScheduledBatchSampler) + for images, targets in loader: + expected_batch_size = {16: 8, 32: 4}[images.shape[-1]] + assert images.shape == (expected_batch_size, 3, images.shape[-2], images.shape[-1]) + assert images.shape[-2] == images.shape[-1] + assert targets.shape == (expected_batch_size,) + + +def test_create_loader_scheduled_resolutions_default_to_batch_size(): + loader = create_loader( + _ImageDataset(), + input_size=(3, 32, 32), + batch_size=4, + input_size_choices=(16, 32), + is_training=True, + no_aug=True, + use_prefetcher=False, + num_workers=0, + persistent_workers=False, + ) + + assert all(images.shape[0] == 4 for images, _ in loader) + + +@pytest.mark.parametrize('invalid_size', (0, (16, 0), (3, 16, 0))) +def test_create_loader_rejects_nonpositive_scheduled_dimensions(invalid_size): + with pytest.raises(ValueError, match='dimensions must be positive'): + create_loader( + _ImageDataset(), + input_size=(3, 32, 32), + batch_size=4, + input_size_choices=(invalid_size, 16), + is_training=True, + no_aug=True, + use_prefetcher=False, + num_workers=0, + persistent_workers=False, + ) + + +def test_create_loader_progressive_resolutions_keep_default_batch_and_sample_counts(): + loader = create_loader( + _ImageDataset(length=96), + input_size=(3, 32, 32), + batch_size=4, + input_size_choices=(16, 24, 32), + batch_choice_schedule='progressive', + batch_schedule_epochs=3, + batch_schedule_spread=0, + batch_schedule_random_mix=0, + is_training=True, + no_aug=True, + use_prefetcher=False, + num_workers=0, + persistent_workers=False, + ) + + assert len(loader) == 24 + for epoch, expected_size in enumerate((16, 24, 32)): + loader.batch_sampler.set_epoch(epoch) + epoch_batches = list(loader) + assert len(epoch_batches) == 24 + assert sum(images.shape[0] for images, _ in epoch_batches) == 96 + assert {images.shape[-1] for images, _ in epoch_batches} == {expected_size} + assert {images.shape[0] for images, _ in epoch_batches} == {4} + + +def test_create_loader_scheduled_resolutions_with_prefetch_mixup(): + loader = create_loader( + _ImageDataset(), + input_size=(3, 32, 32), + batch_size=4, + input_size_choices=(16, 32), + batch_size_choices=(8, 4), + is_training=True, + no_aug=True, + collate_fn=FastCollateMixup(mixup_alpha=0.2, num_classes=64), + use_prefetcher=True, + device=torch.device('cpu'), + num_workers=0, + persistent_workers=False, + ) + + assert isinstance(loader.batch_sampler, ScheduledBatchSampler) + for images, targets in loader: + assert images.dtype == torch.float32 + assert images.shape[0] == targets.shape[0] + assert targets.shape[1] == 64 + + +def test_create_loader_rejects_scheduled_resolutions_with_multi_epochs_loader(): + with pytest.raises(ValueError, match='MultiEpochsDataLoader'): + create_loader( + _ImageDataset(length=32), + input_size=(3, 32, 32), + batch_size=4, + input_size_choices=(16, 32), + is_training=True, + no_aug=True, + use_prefetcher=False, + use_multi_epochs_loader=True, + num_workers=0, + persistent_workers=False, + ) + + +@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_ignores_scheduled_options_without_input_size_choices(scheduled_option): + loader = 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, + ) + + assert loader.batch_size == 4 + + +def test_create_loader_standard_batching_path_is_unchanged(): + dataset = _ImageDataset() + loader = create_loader( + dataset, + input_size=(3, 32, 32), + batch_size=4, + is_training=True, + no_aug=True, + use_prefetcher=False, + num_workers=0, + persistent_workers=False, + ) + + images, targets = next(iter(loader)) + assert loader.dataset is dataset + assert loader.batch_size == 4 + assert images.shape == (4, 3, 32, 32) + assert targets.shape == (4,) diff --git a/timm/data/__init__.py b/timm/data/__init__.py index c54babd48b..bb3ecdcfb8 100644 --- a/timm/data/__init__.py +++ b/timm/data/__init__.py @@ -20,6 +20,7 @@ Patchify, patchify_image, ) +from .scheduled_sampler import ScheduledBatchSampler, ScheduledTransformDataset from .readers import create_reader from .readers import get_img_extensions, is_img_extension, set_img_extensions, add_img_extensions, del_img_extensions from .real_labels import RealLabelsImagenet diff --git a/timm/data/loader.py b/timm/data/loader.py index b6804f6fb5..b13c0190d7 100644 --- a/timm/data/loader.py +++ b/timm/data/loader.py @@ -10,7 +10,7 @@ from contextlib import suppress from functools import partial from itertools import repeat -from typing import Callable, Optional, Tuple, Union +from typing import Callable, Optional, Sequence, Tuple, Union import torch import torch.utils.data @@ -21,6 +21,7 @@ from .distributed_sampler import OrderedDistributedSampler, RepeatAugSampler from .random_erasing import RandomErasing from .mixup import FastCollateMixup +from .scheduled_sampler import ScheduledBatchSampler, ScheduledTransformDataset from .transforms_factory import create_transform _logger = logging.getLogger(__name__) @@ -164,6 +165,10 @@ def __len__(self): def sampler(self): return self.loader.sampler + @property + def batch_sampler(self): + return self.loader.batch_sampler + @property def dataset(self): return self.loader.dataset @@ -237,6 +242,15 @@ def create_loader( persistent_workers: bool = True, worker_seeding: str = 'all', tf_preprocessing: bool = False, + input_size_choices: Optional[Sequence[Union[int, Tuple[int, int], Tuple[int, int, int]]]] = None, + batch_size_choices: Optional[Sequence[int]] = None, + batch_choice_weights: Optional[Sequence[float]] = None, + batch_choice_seed: int = 0, + batch_choice_schedule: str = 'constant', + batch_schedule_epochs: Optional[int] = None, + batch_schedule_spread: float = 0.65, + batch_schedule_random_mix: float = 0.1, + num_batches: Optional[int] = None, ): """ @@ -280,6 +294,18 @@ def create_loader( persistent_workers: Enable persistent worker processes. worker_seeding: Control worker random seeding at init. tf_preprocessing: Use TF 1.0 inference preprocessing for testing model ports. + input_size_choices: Per-batch input size choices for scheduled resolution training. + batch_size_choices: Batch size corresponding to each input size choice. Uses ``batch_size`` for all + choices when not specified. + batch_choice_weights: Sampling weights for input size choices. Uniform when not specified. + batch_choice_seed: Random seed used to create and shuffle the batch schedule. + 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 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. Returns: DataLoader @@ -288,8 +314,7 @@ def create_loader( if re_split: # apply RE to second half of batch if no aug split otherwise line up with aug split re_num_splits = num_aug_splits or 2 - dataset.transform = create_transform( - input_size, + transform_kwargs = dict( is_training=is_training, no_aug=no_aug, train_crop_mode=train_crop_mode, @@ -316,6 +341,53 @@ def create_loader( use_prefetcher=use_prefetcher, separate=num_aug_splits > 0, ) + channels = ( + input_size[0] + if isinstance(input_size, (tuple, list)) and len(input_size) == 3 + else len(mean) + ) + + scheduled_batching = input_size_choices is not None + if scheduled_batching: + if not is_training: + raise ValueError('Scheduled input sizes are only supported for training loaders.') + if use_multi_epochs_loader: + raise ValueError('MultiEpochsDataLoader is not supported with scheduled input sizes.') + if isinstance(dataset, torch.utils.data.IterableDataset): + raise TypeError('Scheduled input sizes require a map-style dataset.') + if num_aug_splits > 0: + raise ValueError('Augmentation splits are not supported with scheduled input sizes.') + if not input_size_choices: + raise ValueError('input_size_choices must contain at least one size.') + + resolved_input_sizes = [] + for size in input_size_choices: + if isinstance(size, int): + size = (channels, size, size) + elif len(size) == 2: + size = (channels, *size) + elif len(size) == 3: + size = tuple(size) + if size[0] != channels: + raise ValueError('All scheduled input sizes must use the same number of channels.') + else: + raise ValueError('Scheduled input sizes must be scalars, HW tuples, or CHW tuples.') + if any(dimension <= 0 for dimension in size): + raise ValueError('All scheduled input size dimensions must be positive.') + resolved_input_sizes.append(size) + + if batch_size_choices is None: + batch_size_choices = [batch_size] * len(resolved_input_sizes) + elif len(batch_size_choices) != len(resolved_input_sizes): + raise ValueError('batch_size_choices and input_size_choices must have the same length.') + if batch_choice_weights is not None and len(batch_choice_weights) != len(resolved_input_sizes): + raise ValueError('batch_choice_weights and input_size_choices must have the same length.') + + transforms = [create_transform(size, **transform_kwargs) for size in resolved_input_sizes] + dataset.transform = None + dataset = ScheduledTransformDataset(dataset, transforms) + else: + dataset.transform = create_transform(input_size, **transform_kwargs) if isinstance(dataset, IterableImageDataset): # give Iterable datasets early knowledge of num_workers so that sample estimates @@ -336,6 +408,9 @@ def create_loader( else: assert num_aug_repeats == 0, "RepeatAugment not currently supported in non-distributed or IterableDataset use" + if scheduled_batching and sampler is None: + sampler = torch.utils.data.RandomSampler(dataset) + if collate_fn is None: collate_fn = fast_collate if use_prefetcher else torch.utils.data.dataloader.default_collate @@ -344,16 +419,32 @@ def create_loader( loader_class = MultiEpochsDataLoader loader_args = dict( - batch_size=batch_size, - shuffle=not isinstance(dataset, torch.utils.data.IterableDataset) and sampler is None and is_training, num_workers=num_workers, - sampler=sampler, collate_fn=collate_fn, pin_memory=pin_memory, - drop_last=is_training, worker_init_fn=partial(_worker_init, worker_seeding=worker_seeding), persistent_workers=persistent_workers ) + if scheduled_batching: + loader_args['batch_sampler'] = ScheduledBatchSampler( + sampler, + batch_sizes=batch_size_choices, + choice_weights=batch_choice_weights, + seed=batch_choice_seed, + drop_last=is_training, + num_batches=num_batches, + choice_schedule=batch_choice_schedule, + schedule_epochs=batch_schedule_epochs, + schedule_spread=batch_schedule_spread, + schedule_random_mix=batch_schedule_random_mix, + ) + else: + loader_args.update( + batch_size=batch_size, + shuffle=not isinstance(dataset, torch.utils.data.IterableDataset) and sampler is None and is_training, + sampler=sampler, + drop_last=is_training, + ) try: loader = loader_class(dataset, **loader_args) except TypeError as e: @@ -365,7 +456,7 @@ def create_loader( loader, mean=mean, std=std, - channels=input_size[0], + channels=channels, device=device, fp16=fp16, # deprecated, use img_dtype img_dtype=img_dtype, diff --git a/timm/data/naflex_dataset.py b/timm/data/naflex_dataset.py index 34a3abf95a..c8c035fadd 100644 --- a/timm/data/naflex_dataset.py +++ b/timm/data/naflex_dataset.py @@ -24,6 +24,7 @@ from PIL import Image from .naflex_transforms import Patchify +from .readers.shared_count import SharedCount from timm.layers import to_2tuple @@ -261,7 +262,7 @@ def __init__( self.distributed = distributed self.rank = rank if distributed else 0 self.world_size = world_size if distributed else 1 - self.epoch = epoch + self.shared_epoch = SharedCount(epoch) self.batch_divisor = batch_divisor # Resolve patch size configuration @@ -303,11 +304,6 @@ def __init__( self._padded_samples_per_rank: int = 0 self._create_canonical_schedule() # Calculate schedule based on padded size - # Per-Epoch State - # Stores (seq_len, list_of_indices) for the current epoch, specific to this rank - self._epoch_batches: List[Tuple[int, List[int]]] = [] - self._prepare_epoch_batches(self.epoch) # setup for initial epoch - def _create_canonical_schedule(self): """ Calculates the canonical batch schedule (seq_len, batch_size pairs) @@ -394,14 +390,17 @@ def _create_canonical_schedule(self): print(f"Rank {self.rank}: Created canonical schedule with {self._num_batches_per_rank} batches for {self._padded_samples_per_rank} samples/rank.") - def _prepare_epoch_batches(self, epoch: int): + def _prepare_epoch_batches(self, epoch: int) -> List[Tuple[int, int, List[int]]]: """ Prepares the batches for the current epoch by: 1. Shuffling the full dataset indices (using epoch seed). 2. Applying padding if in distributed mode. 3. Selecting indices for the current rank. 4. Shuffling the *order* of the canonical batch schedule (using epoch seed). - 5. Assigning the rank's indices to the shuffled batches. + 5. Assigning the rank's indices and patch-size choices to the shuffled batches. + + Returns: + A process-local batch schedule as ``(seq_len, patch_idx, indices)`` tuples. """ g = torch.Generator() g.manual_seed(self.seed + epoch) # Epoch-specific seed for shuffling @@ -454,8 +453,9 @@ def _prepare_epoch_batches(self, epoch: int): else: shuffled_schedule = list(self._canonical_batch_schedule) # Keep original order - # 5. Assign indices to the shuffled batches - self._epoch_batches = [] + # 5. Assign indices and patch-size choices to the shuffled batches + epoch_batches = [] + patch_size_probs = torch.tensor(self.patch_size_probs) idx_pos = 0 scheduled_samples_count = 0 for seq_len, bs in shuffled_schedule: @@ -468,7 +468,10 @@ def _prepare_epoch_batches(self, epoch: int): break # Stop if no more indices or batch size is zero batch_indices = indices_this_rank[idx_pos : idx_pos + actual_bs] - self._epoch_batches.append((seq_len, batch_indices)) + patch_idx = 0 + if self.variable_patch_size: + patch_idx = torch.multinomial(patch_size_probs, 1, generator=g).item() + epoch_batches.append((seq_len, patch_idx, batch_indices)) idx_pos += actual_bs scheduled_samples_count += actual_bs @@ -479,28 +482,26 @@ def _prepare_epoch_batches(self, epoch: int): f"but expected {effective_samples_this_rank} effective samples this epoch. " f"Indices remaining: {effective_samples_this_rank - scheduled_samples_count}." ) + return epoch_batches def set_epoch(self, epoch: int) -> None: - """Updates the epoch, regenerating the epoch-specific batches. + """Set the multiprocessing-safe epoch read by DataLoader workers. Args: epoch: New epoch number. """ - # Only regenerate if the epoch actually changes - if epoch != self.epoch: - self.epoch = epoch - self._prepare_epoch_batches(epoch) + self.shared_epoch.value = epoch def __len__(self) -> int: - """Returns the number of batches per worker for the current epoch. + """Return the number of batches for this rank. Returns: - Number of batches this worker will process. + Number of batches collectively processed by this rank's workers. """ return self._num_batches_per_rank def __iter__(self) -> Iterator[Tuple[Dict[str, torch.Tensor], torch.Tensor]]: - """Iterates through pre-calculated batches for the current epoch. + """Prepare and iterate this worker's batches for the shared epoch. Yields: Tuple of (input_dict, targets) for each batch. @@ -509,19 +510,19 @@ def __iter__(self) -> Iterator[Tuple[Dict[str, torch.Tensor], torch.Tensor]]: num_workers = worker_info.num_workers if worker_info else 1 worker_id = worker_info.id if worker_info else 0 - # Distribute pre-calculated batches among workers for this rank - # Each worker processes a slice of the batches prepared in _prepare_epoch_batches - batches_for_worker = self._epoch_batches[worker_id::num_workers] - for seq_len, indices in batches_for_worker: + # Snapshot the shared epoch once, then build process-local derived state. + epoch = self.shared_epoch.value + 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 - # Select patch size for this batch - patch_idx = 0 - if self.variable_patch_size: - # Use torch multinomial for weighted random choice - patch_idx = torch.multinomial(torch.tensor(self.patch_size_probs), 1).item() - # Get the pre-initialized transform and patchifier using patch_idx transform_key = (seq_len, patch_idx) transform = self.transforms.get(transform_key) diff --git a/timm/data/naflex_loader.py b/timm/data/naflex_loader.py index 252f36d6d7..735decdfe5 100644 --- a/timm/data/naflex_loader.py +++ b/timm/data/naflex_loader.py @@ -207,6 +207,11 @@ def sampler(self): """ return self.loader.sampler + @property + def batch_sampler(self): + """Get batch sampler from underlying loader.""" + return self.loader.batch_sampler + @property def dataset(self): """Get dataset from underlying loader. diff --git a/timm/data/scheduled_sampler.py b/timm/data/scheduled_sampler.py new file mode 100644 index 0000000000..8474e50cb5 --- /dev/null +++ b/timm/data/scheduled_sampler.py @@ -0,0 +1,331 @@ +"""Scheduled batch sampling and transform dispatch for map-style datasets.""" + +import math +from itertools import islice +from typing import Any, Callable, Iterator, List, Optional, Sequence, Tuple, Union + +import torch +from torch.utils.data import Dataset, Sampler + + +class ScheduledBatchSampler(Sampler[List[Tuple[Any, int]]]): + """Batch sampler that associates each batch with a scheduled transform choice. + + In sample-budget mode, a canonical composition of ``(choice index, batch + size)`` pairs is created once and shuffled at the start of each iteration. + Fixed-batch progressive mode creates a schedule for the current epoch, + moving a probability window from the first choice to the last while + retaining configurable random exploration. + + The same seed, sampler length, and epoch produce an identical batch-shape + schedule on every distributed rank. If a fixed number of batches requires + more indices than the wrapped sampler provides, its index stream is cycled. + + The yielded indices are ``(sample index, choice index)`` pairs intended for + use with :class:`ScheduledTransformDataset`. + + Args: + sampler: Base sample-index sampler. + batch_sizes: Batch size corresponding to each transform choice. + choice_weights: Base sampling weights for the choices. A zero weight + excludes that choice from all schedules. + seed: Seed shared by schedule creation and shuffling. + drop_last: Drop an undersized tail in sample-budget mode. + shuffle_schedule: Shuffle batch order within each epoch. + num_batches: Fixed number of batches per epoch. Progressive mode infers + this from its training-average expected batch size when omitted. + choice_schedule: Either ``constant`` or ``progressive``. + schedule_epochs: Number of epochs over which progressive choices move + from the first choice to the last. + schedule_spread: Progressive probability-window standard deviation in + choice-index units. Zero selects only the nearest choice. + schedule_random_mix: Fraction of uniform exploration over active + choices mixed into the progressive choice probabilities. + """ + + def __init__( + self, + sampler: Sampler, + batch_sizes: Sequence[int], + choice_weights: Optional[Sequence[float]] = None, + seed: int = 0, + drop_last: bool = True, + shuffle_schedule: bool = True, + num_batches: Optional[int] = None, + choice_schedule: str = 'constant', + schedule_epochs: Optional[int] = None, + schedule_spread: float = 0.65, + schedule_random_mix: float = 0.1, + ) -> None: + if not hasattr(sampler, '__len__'): + raise TypeError('ScheduledBatchSampler requires a sampler with a length.') + if len(sampler) <= 0: + raise ValueError('ScheduledBatchSampler requires a non-empty sampler.') + if not batch_sizes: + raise ValueError('batch_sizes must contain at least one value.') + if any(int(batch_size) != batch_size or batch_size <= 0 for batch_size in batch_sizes): + raise ValueError('All scheduled batch sizes must be positive integers.') + if num_batches is not None and (int(num_batches) != num_batches or num_batches <= 0): + raise ValueError('num_batches must be a positive integer when specified.') + if choice_schedule not in ('constant', 'progressive'): + raise ValueError("choice_schedule must be 'constant' or 'progressive'.") + if choice_schedule == 'progressive': + if len(batch_sizes) < 2: + raise ValueError('A progressive schedule requires at least two choices.') + if schedule_epochs is None or int(schedule_epochs) != schedule_epochs or schedule_epochs <= 0: + raise ValueError('schedule_epochs must be a positive integer for a progressive schedule.') + if schedule_spread < 0: + raise ValueError('schedule_spread must be non-negative.') + if not 0 <= schedule_random_mix <= 1: + raise ValueError('schedule_random_mix must be between 0 and 1.') + + self.sampler = sampler + self.batch_sizes = tuple(int(batch_size) for batch_size in batch_sizes) + self.choice_weights = self._normalize_choice_weights(choice_weights) + self._active_choices = tuple( + choice_index + for choice_index, choice_weight in enumerate(self.choice_weights) + if choice_weight > 0 + ) + self.seed = seed + self.drop_last = drop_last + self.shuffle_schedule = shuffle_schedule + self.choice_schedule = choice_schedule + self.schedule_epochs = int(schedule_epochs) if schedule_epochs is not None else None + self.schedule_spread = schedule_spread + self.schedule_random_mix = schedule_random_mix + self.epoch = 0 + self.average_batch_size = self._calculate_average_batch_size() + if choice_schedule == 'progressive' and num_batches is None: + num_batches = self._infer_num_batches() + self.num_batches = int(num_batches) if num_batches is not None else None + self._sample_budget_schedule: Tuple[Tuple[int, int], ...] = () + if self.num_batches is None: + self._sample_budget_schedule = self._create_sample_budget_schedule() + if not self._sample_budget_schedule: + raise ValueError( + 'No full scheduled batch fits the sampler; reduce the batch sizes.' + ) + + def _normalize_choice_weights( + self, + choice_weights: Optional[Sequence[float]], + ) -> torch.Tensor: + if choice_weights is None: + return torch.full((len(self.batch_sizes),), 1.0 / len(self.batch_sizes), dtype=torch.float64) + if len(choice_weights) != len(self.batch_sizes): + raise ValueError('choice_weights and batch_sizes must have the same length.') + + weights = torch.tensor(choice_weights, dtype=torch.float64) + if not torch.isfinite(weights).all() or (weights < 0).any(): + raise ValueError('choice_weights must contain finite, non-negative values.') + weight_sum = weights.sum() + if weight_sum <= 0: + raise ValueError('choice_weights must have a positive sum.') + return weights / weight_sum + + def choice_weights_for_epoch(self, epoch: int) -> torch.Tensor: + """Return normalized choice weights for an epoch. + + Args: + epoch: Zero-based training epoch. + + Returns: + Normalized floating-point choice weights. + """ + if self.choice_schedule == 'constant' or len(self.batch_sizes) == 1: + return self.choice_weights + + if self.schedule_epochs == 1: + progress = 1.0 + else: + progress = min(max(epoch / (self.schedule_epochs - 1), 0.0), 1.0) + choice_positions = torch.arange(len(self.batch_sizes), dtype=torch.float64) + center = progress * (len(self.batch_sizes) - 1) + + if self.schedule_spread == 0: + distances = (choice_positions - center).abs() + curriculum_weights = (distances == distances.min()).to(torch.float64) + else: + curriculum_weights = torch.exp( + -0.5 * ((choice_positions - center) / self.schedule_spread) ** 2 + ) + curriculum_weights *= self.choice_weights + if curriculum_weights.sum() <= 0: + nearest_active = min(self._active_choices, key=lambda index: abs(index - center)) + curriculum_weights = torch.zeros_like(self.choice_weights) + curriculum_weights[nearest_active] = 1.0 + else: + curriculum_weights /= curriculum_weights.sum() + + if self.schedule_random_mix: + random_weights = (self.choice_weights > 0).to(curriculum_weights.dtype) + random_weights /= random_weights.sum() + curriculum_weights = ( + 1.0 - self.schedule_random_mix + ) * curriculum_weights + self.schedule_random_mix * random_weights + return curriculum_weights / curriculum_weights.sum() + + def _calculate_average_batch_size(self) -> float: + batch_sizes = torch.tensor(self.batch_sizes, dtype=torch.float64) + if self.choice_schedule == 'progressive': + expected_batch_sizes = [ + torch.dot(self.choice_weights_for_epoch(epoch), batch_sizes) for epoch in range(self.schedule_epochs) + ] + return float(torch.stack(expected_batch_sizes).mean().item()) + return float(torch.dot(self.choice_weights, batch_sizes).item()) + + def _infer_num_batches(self) -> int: + if len(set(self.batch_sizes)) == 1: + batch_size = self.batch_sizes[0] + if self.drop_last: + num_batches = len(self.sampler) // batch_size + else: + num_batches = math.ceil(len(self.sampler) / batch_size) + else: + if self.drop_last: + num_batches = int(len(self.sampler) / self.average_batch_size) + else: + num_batches = math.ceil(len(self.sampler) / self.average_batch_size) + if num_batches < 1: + raise ValueError('No full scheduled batch fits the sampler; reduce the batch sizes.') + return num_batches + + def _sample_choice( + self, + generator: torch.Generator, + valid_choices: Optional[Sequence[int]] = None, + choice_weights: Optional[torch.Tensor] = None, + ) -> int: + choice_weights = self.choice_weights if choice_weights is None else choice_weights + if valid_choices is None: + return int(torch.multinomial(choice_weights, 1, generator=generator).item()) + + valid_choices = tuple(valid_choices) + weights = choice_weights[list(valid_choices)] + if weights.sum() <= 0: + raise RuntimeError('No positive-weight scheduled choice is available for this batch.') + sampled_index = int(torch.multinomial(weights, 1, generator=generator).item()) + return valid_choices[sampled_index] + + def _create_sample_budget_schedule(self) -> Tuple[Tuple[int, int], ...]: + generator = torch.Generator().manual_seed(self.seed) + remaining = len(self.sampler) + min_batch_size = min(self.batch_sizes[index] for index in self._active_choices) + schedule = [] + + while remaining >= min_batch_size: + valid_choices = [ + choice_index for choice_index in self._active_choices if self.batch_sizes[choice_index] <= remaining + ] + choice_index = self._sample_choice(generator, valid_choices) + batch_size = self.batch_sizes[choice_index] + schedule.append((choice_index, batch_size)) + remaining -= batch_size + + if remaining and not self.drop_last: + schedule.append((self._sample_choice(generator), remaining)) + + return tuple(schedule) + + def _create_fixed_batch_schedule(self, epoch: int) -> Tuple[Tuple[int, int], ...]: + generator = torch.Generator().manual_seed(self.seed + 2 * epoch) + choice_weights = self.choice_weights_for_epoch(epoch) + schedule = [] + for _ in range(self.num_batches): + choice_index = self._sample_choice(generator, choice_weights=choice_weights) + schedule.append((choice_index, self.batch_sizes[choice_index])) + return tuple(schedule) + + def _create_schedule(self, epoch: int) -> Tuple[Tuple[int, int], ...]: + if self.num_batches is not None: + return self._create_fixed_batch_schedule(epoch) + return self._sample_budget_schedule + + @property + def schedule(self) -> Tuple[Tuple[int, int], ...]: + """Return the unshuffled schedule for the currently selected epoch.""" + return self._create_schedule(self.epoch) + + def _cycling_sampler(self) -> Iterator[Any]: + while True: + yielded = False + for sample_index in self.sampler: + yielded = True + yield sample_index + if not yielded: + break + + def __iter__(self) -> Iterator[List[Tuple[Any, int]]]: + epoch = self.epoch + schedule = self._create_schedule(epoch) + if self.shuffle_schedule and len(schedule) > 1: + generator = torch.Generator().manual_seed(self.seed + 2 * epoch + 1) + order = torch.randperm(len(schedule), generator=generator).tolist() + schedule = tuple(schedule[index] for index in order) + + sample_indices = self._cycling_sampler() if self.num_batches is not None else iter(self.sampler) + for choice_index, batch_size in schedule: + batch = list(islice(sample_indices, batch_size)) + if len(batch) != batch_size: + if self.drop_last or not batch: + break + yield [(sample_index, choice_index) for sample_index in batch] + + def __len__(self) -> int: + if self.num_batches is not None: + return self.num_batches + return len(self._sample_budget_schedule) + + def set_epoch(self, epoch: int) -> None: + """Set the epoch for the next iteration and the wrapped sampler.""" + self.epoch = epoch + if hasattr(self.sampler, 'set_epoch'): + self.sampler.set_epoch(epoch) + + +class ScheduledTransformDataset(Dataset): + """Apply one of several transforms selected by an annotated sample index. + + Args: + dataset: Map-style base dataset returning tuple or list samples. + transforms: Transforms indexed by the choice annotation supplied by + :class:`ScheduledBatchSampler`. + """ + + def __init__( + self, + dataset: Dataset, + transforms: Sequence[Callable], + ) -> None: + if not transforms: + raise ValueError('transforms must contain at least one transform.') + self.dataset = dataset + self.transforms = tuple(transforms) + + def __getitem__(self, scheduled_index: Tuple[Any, int]) -> Union[Tuple[Any, ...], List[Any]]: + sample_index, transform_index = scheduled_index + if not 0 <= transform_index < len(self.transforms): + raise IndexError(f'Transform index {transform_index} is out of range.') + + sample = self.dataset[sample_index] + if not isinstance(sample, (tuple, list)) or not sample: + raise TypeError('ScheduledTransformDataset expects tuple/list dataset samples.') + image = self.transforms[transform_index](sample[0]) + if isinstance(sample, tuple): + return image, *sample[1:] + return [image, *sample[1:]] + + def __len__(self) -> int: + return len(self.dataset) + + def set_epoch(self, epoch: int) -> None: + """Forward epoch updates to an epoch-aware base dataset.""" + if hasattr(self.dataset, 'set_epoch'): + self.dataset.set_epoch(epoch) + + def filename(self, index: int, basename: bool = False, absolute: bool = False) -> Any: + return self.dataset.filename(index, basename=basename, absolute=absolute) + + def filenames(self, basename: bool = False, absolute: bool = False) -> Any: + return self.dataset.filenames(basename=basename, absolute=absolute) diff --git a/train.py b/train.py index 74acfd6d2b..4b9f3082f4 100755 --- a/train.py +++ b/train.py @@ -401,6 +401,24 @@ group.add_argument('--wandb-resume-id', default='', type=str, metavar='ID', help='If resuming a run, the id of the run in wandb') +# Scheduled resolution loader arguments +group.add_argument('--train-img-sizes', type=int, nargs='+', default=None, + help='Per-batch square image sizes for scheduled resolution training') +group.add_argument('--train-batch-sizes', type=int, nargs='+', default=None, + help='Batch size for each --train-img-sizes choice (default: --batch-size for every choice)') +group.add_argument('--train-size-probs', type=float, nargs='+', default=None, + help='Base sampling weights for --train-img-sizes (default: uniform)') +group.add_argument('--train-size-schedule', type=str, default='constant', choices=('constant', 'progressive'), + help='Resolution choice schedule; progressive moves from the first size to the last') +group.add_argument('--train-size-schedule-spread', type=float, default=0.65, + help='Progressive schedule spread in resolution-choice index units (default: 0.65)') +group.add_argument('--train-size-random-mix', type=float, default=0.1, + help='Fraction of uniform random active choices mixed into a progressive schedule (default: 0.1)') +group.add_argument('--train-batches-per-epoch', '--train-steps-per-epoch', type=int, default=None, + help='Fixed loader batches (before gradient accumulation) per epoch; inferred if unspecified') +group.add_argument('--variable-batch-loss-scale', default='none', type=str, choices=('none', 'sqrt', 'linear'), + help='Scale gradients relative to the policy-average scheduled batch size') + # NaFlex scheduled loader arguments group.add_argument('--naflex-loader', action='store_true', default=False, help='Use NaFlex loader (Requires NaFlex compatible model)') @@ -412,7 +430,7 @@ help='List of patch sizes for variable patch size training (e.g., 8 12 16 24 32)') group.add_argument('--naflex-patch-size-probs', type=float, nargs='+', default=None, help='Probabilities for each patch size (must sum to 1.0, uniform if not specified)') -group.add_argument('--naflex-loss-scale', default='linear', type=str, +group.add_argument('--naflex-loss-scale', default='linear', type=str, choices=('none', 'sqrt', 'linear'), help='Scale loss (gradient) by batch_size ("none", "sqrt", or "linear")') group.add_argument('--naflex-patchify-channels-first', action='store_true', default=False, help='Emit NaFlex patches in C-P-P (channels-first) flat layout instead of the default ' @@ -457,6 +475,15 @@ def _parse_args(): return args, args_text +def _set_loader_epoch(loader, epoch: int) -> None: + if hasattr(loader.dataset, 'set_epoch'): + loader.dataset.set_epoch(epoch) + if hasattr(loader.batch_sampler, 'set_epoch'): + loader.batch_sampler.set_epoch(epoch) + elif hasattr(loader.sampler, 'set_epoch'): + loader.sampler.set_epoch(epoch) + + def main(): utils.setup_default_logging() args, args_text = _parse_args() @@ -471,6 +498,10 @@ def main(): args.prefetcher = not args.no_prefetcher args.grad_accum_steps = max(1, args.grad_accum_steps) + if args.naflex_loader and args.train_img_sizes is not None: + parser.error('--naflex-loader and --train-img-sizes are alternative loader modes.') + if args.naflex_loader and args.use_multi_epochs_loader: + parser.error('--use-multi-epochs-loader is not supported with --naflex-loader.') device = utils.init_distributed_device(args) if args.distributed: _logger.info( @@ -593,20 +624,6 @@ def main(): assert not args.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model' model = torch.jit.script(model) - if not args.lr: - global_batch_size = args.batch_size * args.world_size * args.grad_accum_steps - batch_ratio = global_batch_size / args.lr_base_size - if not args.lr_base_scale: - on = args.opt.lower() - args.lr_base_scale = 'sqrt' if any([o in on for o in ('ada', 'lamb')]) else 'linear' - if args.lr_base_scale == 'sqrt': - batch_ratio = batch_ratio ** 0.5 - args.lr = args.lr_base * batch_ratio - if utils.is_primary(args): - _logger.info( - f'Learning rate ({args.lr}) calculated from base learning rate ({args.lr_base}) ' - f'and effective global batch size ({global_batch_size}) with {args.lr_base_scale} scaling.') - # setup automatic mixed-precision (AMP) loss scaling and op casting amp_autocast = suppress # do nothing loss_scaler = None @@ -720,6 +737,8 @@ def main(): ) naflex_mode = False + scheduled_batch_mode = args.train_img_sizes is not None + batch_size_reference = float(args.batch_size) if args.naflex_loader: if utils.is_primary(args): _logger.info('Using NaFlex loader') @@ -786,11 +805,48 @@ def main(): loader_train = create_loader( dataset_train, input_size=data_config['input_size'], + input_size_choices=args.train_img_sizes, + batch_size_choices=args.train_batch_sizes, + batch_choice_weights=args.train_size_probs, + batch_choice_seed=args.seed, + batch_choice_schedule=args.train_size_schedule, + batch_schedule_epochs=args.epochs if args.train_size_schedule == 'progressive' else None, + batch_schedule_spread=args.train_size_schedule_spread, + batch_schedule_random_mix=args.train_size_random_mix, + num_batches=args.train_batches_per_epoch, collate_fn=collate_fn, use_multi_epochs_loader=args.use_multi_epochs_loader, **common_loader_kwargs, **train_loader_kwargs, ) + if scheduled_batch_mode: + scheduled_sampler = loader_train.batch_sampler + batch_size_reference = scheduled_sampler.average_batch_size + if utils.is_primary(args): + _logger.info( + f'Using scheduled training resolutions {args.train_img_sizes} with batch sizes ' + f'{scheduled_sampler.batch_sizes}.') + if args.train_size_schedule == 'progressive': + _logger.info( + f'Using progressive resolution schedule over {args.epochs} epochs with ' + f'spread={args.train_size_schedule_spread} and random_mix={args.train_size_random_mix}.') + _logger.info( + f'Scheduled loader has {len(scheduled_sampler)} batches/epoch using ' + f'policy-average batch size {batch_size_reference:.2f}.') + + if not args.lr: + global_batch_size = batch_size_reference * args.world_size * args.grad_accum_steps + batch_ratio = global_batch_size / args.lr_base_size + if not args.lr_base_scale: + on = args.opt.lower() + args.lr_base_scale = 'sqrt' if any([o in on for o in ('ada', 'lamb')]) else 'linear' + if args.lr_base_scale == 'sqrt': + batch_ratio = batch_ratio ** 0.5 + args.lr = args.lr_base * batch_ratio + if utils.is_primary(args): + _logger.info( + f'Learning rate ({args.lr}) calculated from base learning rate ({args.lr_base}) ' + f'and effective global batch size ({global_batch_size:g}) with {args.lr_base_scale} scaling.') loader_eval = None if args.val_split: @@ -1050,10 +1106,7 @@ def main(): results = [] try: for epoch in range(start_epoch, num_epochs): - if hasattr(dataset_train, 'set_epoch'): - dataset_train.set_epoch(epoch) - elif args.distributed and hasattr(loader_train.sampler, 'set_epoch'): - loader_train.sampler.set_epoch(epoch) + _set_loader_epoch(loader_train, epoch) train_metrics = train_one_epoch( epoch, @@ -1072,6 +1125,8 @@ def main(): mixup_fn=mixup_fn, num_updates_total=num_epochs * updates_per_epoch, naflex_mode=naflex_mode, + scheduled_batch_mode=scheduled_batch_mode, + batch_size_reference=batch_size_reference, ) if args.distributed and args.dist_bn in ('broadcast', 'reduce'): @@ -1190,6 +1245,8 @@ def train_one_epoch( mixup_fn=None, num_updates_total=None, naflex_mode=False, + scheduled_batch_mode=False, + batch_size_reference=None, ): if args.mixup_off_epoch and epoch >= args.mixup_off_epoch: if args.prefetcher and loader.mixup_enabled: @@ -1271,16 +1328,21 @@ def _backward(_loss): if naflex_mode: assert isinstance(input, dict) batch_size = input['patches'].shape[0] + else: + batch_size = input.shape[0] - # scale gradient vs the minimum batch size (for max seq len) - if not args.naflex_loss_scale or args.naflex_loss_scale == 'none': + if naflex_mode or scheduled_batch_mode: + loss_scale_mode = args.naflex_loss_scale if naflex_mode else args.variable_batch_loss_scale + if not loss_scale_mode or loss_scale_mode == 'none': local_scale = 1.0 else: - local_scale = (batch_size / args.batch_size) - if local_scale == 'sqrt': + local_scale = batch_size / batch_size_reference + if loss_scale_mode == 'sqrt': local_scale = local_scale ** 0.5 + elif loss_scale_mode != 'linear': + raise ValueError(f'Invalid variable batch loss scale mode: {loss_scale_mode}') - if args.distributed: + if naflex_mode and args.distributed: # scale gradient btw distributed ranks, each one can have different batch size global_batch_size = utils.reduce_tensor( torch.tensor(batch_size, device=device, dtype=torch.float32), @@ -1290,6 +1352,8 @@ def _backward(_loss): else: dist_scale = None global_batch_size = batch_size + if args.distributed: + global_batch_size *= args.world_size if has_no_sync and not need_update: with task.no_sync(): @@ -1305,7 +1369,7 @@ def _backward(_loss): scaled_loss *= dist_scale _backward(scaled_loss) else: - global_batch_size = batch_size = input.shape[0] + global_batch_size = batch_size if args.distributed: global_batch_size *= args.world_size