diff --git a/examples/imagenet-example b/examples/imagenet-example index d394723e..f134cbff 160000 --- a/examples/imagenet-example +++ b/examples/imagenet-example @@ -1 +1 @@ -Subproject commit d394723e41e023017562df86a0e95c04be5cd119 +Subproject commit f134cbfff7f590954edc5c24275444b7dd2f57f6 diff --git a/ffcv/__init__.py b/ffcv/__init__.py index cbbcbd22..7fdfb1d6 100644 --- a/ffcv/__init__.py +++ b/ffcv/__init__.py @@ -1,5 +1,5 @@ from .loader import Loader from .writer import DatasetWriter -__version__ = '0.0.2' +__version__ = '0.0.3rc1' __all__ = ['Loader'] diff --git a/ffcv/loader/epoch_iterator.py b/ffcv/loader/epoch_iterator.py index 4bc1b9bd..b54fa96b 100644 --- a/ffcv/loader/epoch_iterator.py +++ b/ffcv/loader/epoch_iterator.py @@ -46,9 +46,8 @@ def __init__(self, loader: 'Loader', order: Sequence[int]): self.memory_bank_per_stage = defaultdict(list) - if IS_CUDA: - self.cuda_streams = [ch.cuda.Stream() - for _ in range(self.loader.batches_ahead + 2)] + self.cuda_streams = [(ch.cuda.Stream() if IS_CUDA else None) + for _ in range(self.loader.batches_ahead + 2)] # Allocate all the memory memory_allocations = {} @@ -136,7 +135,7 @@ def run_pipeline(self, b_ix, batch_indices, batch_slot, cuda_event): if first_stage: first_stage = False self.memory_context.end_batch(b_ix) - return tuple(args) + return tuple(x[:len(batch_indices)] for x in args) def __next__(self): result = self.output_queue.get() diff --git a/ffcv/loader/loader.py b/ffcv/loader/loader.py index 4fcc8678..21cdd104 100644 --- a/ffcv/loader/loader.py +++ b/ffcv/loader/loader.py @@ -93,7 +93,7 @@ def __init__(self, os_cache: bool = DEFAULT_OS_CACHE, order: ORDER_TYPE = OrderOption.SEQUENTIAL, distributed: bool = False, - seed: int = 0, # For ordering of samples + seed: int = None, # For ordering of samples indices: Sequence[int] = None, # For subset selection pipelines: Mapping[str, Sequence[Union[Operation, ch.nn.Module]]] = {}, @@ -103,6 +103,14 @@ def __init__(self, recompile: bool = False, # Recompile at every epoch ): + if distributed and order == OrderOption.RANDOM and (seed is None): + print('Warning: no ordering seed was specified with distributed=True. ' + 'Setting seed to 0 to match PyTorch distributed sampler.') + seed = 0 + elif seed is None: + tinfo = np.iinfo('int32') + seed = np.random.randint(0, tinfo.max) + # We store the original user arguments to be able to pass it to the # filtered version of the datasets self._args = { diff --git a/ffcv/transforms/__init__.py b/ffcv/transforms/__init__.py index 0850d146..bc8fa321 100644 --- a/ffcv/transforms/__init__.py +++ b/ffcv/transforms/__init__.py @@ -7,6 +7,8 @@ from .replace_label import ReplaceLabel from .normalize import NormalizeImage from .translate import RandomTranslate +from .mixup import ImageMixup, LabelMixup, MixupToOneHot +from .module import ModuleWrapper __all__ = ['ToTensor', 'ToDevice', 'ToTorchImage', 'NormalizeImage', diff --git a/ffcv/transforms/cutout.py b/ffcv/transforms/cutout.py index a7402bc5..89237e0e 100644 --- a/ffcv/transforms/cutout.py +++ b/ffcv/transforms/cutout.py @@ -3,6 +3,7 @@ """ import numpy as np from typing import Callable, Optional, Tuple +from dataclasses import replace from ffcv.pipeline.compiler import Compiler from ..pipeline.allocation_query import AllocationQuery @@ -48,5 +49,4 @@ def cutout_square(images, *_): return cutout_square def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: - assert previous_state.jit_mode - return previous_state, None + return replace(previous_state, jit_mode=True), None diff --git a/ffcv/transforms/flip.py b/ffcv/transforms/flip.py index b0296f1a..63d4b1f9 100644 --- a/ffcv/transforms/flip.py +++ b/ffcv/transforms/flip.py @@ -1,7 +1,7 @@ """ Random horizontal flip """ -from numpy import dtype +from dataclasses import replace from numpy.random import rand from typing import Callable, Optional, Tuple from ..pipeline.allocation_query import AllocationQuery @@ -42,5 +42,5 @@ def flip(images, dst): return flip def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: - assert previous_state.jit_mode - return (previous_state, AllocationQuery(previous_state.shape, previous_state.dtype)) + return (replace(previous_state, jit_mode=True), + AllocationQuery(previous_state.shape, previous_state.dtype)) diff --git a/ffcv/transforms/mixup.py b/ffcv/transforms/mixup.py index 25741001..53239b6f 100644 --- a/ffcv/transforms/mixup.py +++ b/ffcv/transforms/mixup.py @@ -53,8 +53,6 @@ def mixer(images, dst, indices): return mixer def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: - # assert previous_state.jit_mode - # We do everything in place return (previous_state, AllocationQuery(shape=previous_state.shape, dtype=previous_state.dtype)) @@ -92,8 +90,6 @@ def mixer(labels, temp_array, indices): return mixer def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: - # assert previous_state.jit_mode - # We do everything in place return (replace(previous_state, shape=(3,), dtype=np.float32), AllocationQuery((3,), dtype=np.float32)) @@ -115,6 +111,7 @@ def one_hotter(mixedup_labels, dst): return one_hotter def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: + # Should already be converted to tensor assert not previous_state.jit_mode return (replace(previous_state, shape=(self.num_classes,)), \ AllocationQuery((self.num_classes,), dtype=previous_state.dtype, device=previous_state.device)) \ No newline at end of file diff --git a/ffcv/transforms/poisoning.py b/ffcv/transforms/poisoning.py index 7df9ca8d..6b897885 100644 --- a/ffcv/transforms/poisoning.py +++ b/ffcv/transforms/poisoning.py @@ -1,13 +1,10 @@ """ Poison images by adding a mask """ -from collections.abc import Sequence from typing import Tuple +from dataclasses import replace import numpy as np -from numpy import dtype -from numpy.core.numeric import indices -from numpy.random import rand from typing import Callable, Optional, Tuple from ..pipeline.allocation_query import AllocationQuery from ..pipeline.operation import Operation @@ -67,6 +64,6 @@ def poison(images, temp_array, indices): return poison def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: - assert previous_state.jit_mode # We do everything in place - return (previous_state, AllocationQuery(shape=previous_state.shape, dtype=np.float32)) + return (replace(previous_state, jit_mode=True), \ + AllocationQuery(shape=previous_state.shape, dtype=np.dtype('float32'))) diff --git a/ffcv/transforms/replace_label.py b/ffcv/transforms/replace_label.py index e69ec9e9..5a95c011 100644 --- a/ffcv/transforms/replace_label.py +++ b/ffcv/transforms/replace_label.py @@ -1,13 +1,10 @@ """ Replace label """ -from collections.abc import Sequence from typing import Tuple import numpy as np -from numpy import dtype -from numpy.core.numeric import indices -from numpy.random import rand +from dataclasses import replace from typing import Callable, Optional, Tuple from ..pipeline.allocation_query import AllocationQuery from ..pipeline.operation import Operation @@ -50,6 +47,4 @@ def replace_label(labels, temp_array, indices): return replace_label def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: - assert previous_state.jit_mode - # We do everything in place - return (previous_state, None) + return (replace(previous_state, jit_mode=True), None) diff --git a/ffcv/transforms/translate.py b/ffcv/transforms/translate.py index bb2ab5d5..e53e157e 100644 --- a/ffcv/transforms/translate.py +++ b/ffcv/transforms/translate.py @@ -2,9 +2,9 @@ Random translate """ import numpy as np -from numpy import dtype from numpy.random import randint -from typing import Any, Callable, Optional, Tuple, Union +from typing import Callable, Optional, Tuple +from dataclasses import replace from ..pipeline.allocation_query import AllocationQuery from ..pipeline.operation import Operation from ..pipeline.state import State @@ -51,5 +51,6 @@ def translate(images, dst): def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: h, w, c = previous_state.shape - assert previous_state.jit_mode - return (previous_state, AllocationQuery((h + 2 * self.padding, w + 2 * self.padding, c), previous_state.dtype)) + return (replace(previous_state, jit_mode=True), \ + AllocationQuery((h + 2 * self.padding, w + 2 * self.padding, c), previous_state.dtype)) + diff --git a/ffcv/traversal_order/random.py b/ffcv/traversal_order/random.py index cef4ac24..25bc1708 100644 --- a/ffcv/traversal_order/random.py +++ b/ffcv/traversal_order/random.py @@ -24,4 +24,4 @@ def sample_order(self, epoch: int) -> Sequence[int]: self.sampler.set_epoch(epoch) - return np.array(list(self.sampler)) + return self.indices[np.array(list(self.sampler))] diff --git a/ffcv/traversal_order/sequential.py b/ffcv/traversal_order/sequential.py index 632a0059..c5bba224 100644 --- a/ffcv/traversal_order/sequential.py +++ b/ffcv/traversal_order/sequential.py @@ -27,4 +27,4 @@ def sample_order(self, epoch: int) -> Sequence[int]: self.sampler.set_epoch(epoch) - return np.array(list(self.sampler)) + return self.indices[np.array(list(self.sampler))] diff --git a/setup.py b/setup.py index 075765d5..978a6df9 100644 --- a/setup.py +++ b/setup.py @@ -37,11 +37,12 @@ def pkgconfig(package, kw): **extension_kwargs) setup(name='ffcv', - version='0.0.2', + version='0.0.3rc1', description=' FFCV: Fast Forward Computer Vision ', author='MadryLab', author_email='leclerc@mit.edu', - url='https://github.com/MadryLab/fastercv', + url='https://github.com/libffcv/ffcv', + license_files = ('LICENSE.txt',), packages=find_packages(), long_description=long_description, long_description_content_type='text/markdown', diff --git a/tests/test_augmentations.py b/tests/test_augmentations.py index 0d8776b3..f4a44530 100644 --- a/tests/test_augmentations.py +++ b/tests/test_augmentations.py @@ -1,21 +1,35 @@ +import os +import uuid import numpy as np import torch as ch from torch.utils.data import Dataset +from torchvision import transforms as tvt from assertpy import assert_that from tempfile import NamedTemporaryFile from torchvision.datasets import CIFAR10 +from torchvision.utils import save_image, make_grid from torch.utils.data import Subset from ffcv.fields.basics import IntDecoder from ffcv.fields.rgb_image import SimpleRGBImageDecoder -from ffcv.transforms.cutout import Cutout from ffcv.writer import DatasetWriter from ffcv.fields import IntField, RGBImageField from ffcv.loader import Loader from ffcv.pipeline.compiler import Compiler -from ffcv.transforms import Squeeze, Cutout, ToTensor, ToDevice, Poison +from ffcv.transforms import * -def run_test(length, pipeline, compile): +SAVE_IMAGES = True +IMAGES_TMP_PATH = '/tmp/ffcv_augtest_output' +if SAVE_IMAGES: + os.makedirs(IMAGES_TMP_PATH, exist_ok=True) + +UNAUGMENTED_PIPELINE=[ + SimpleRGBImageDecoder(), + ToTensor(), + ToTorchImage() +] + +def run_test(length, pipeline, compile=False): my_dataset = Subset(CIFAR10(root='/tmp', train=True, download=True), range(length)) with NamedTemporaryFile() as handle: @@ -28,33 +42,75 @@ def run_test(length, pipeline, compile): writer.from_indexed_dataset(my_dataset, chunksize=10) - Compiler.set_enabled(True) + Compiler.set_enabled(compile) loader = Loader(name, batch_size=7, num_workers=2, pipelines={ 'image': pipeline, 'label': [IntDecoder(), ToTensor(), Squeeze()] }, drop_last=False) + + unaugmented_loader = Loader(name, batch_size=7, num_workers=2, pipelines={ + 'image': UNAUGMENTED_PIPELINE, + 'label': [IntDecoder(), ToTensor(), Squeeze()] + }, drop_last=False) + tot_indices = 0 tot_images = 0 - for images, label in loader: - tot_indices += label.shape[0] + for (images, labels), (original_images, original_labels) in zip(loader, unaugmented_loader): + print(images.shape, original_images.shape) + tot_indices += labels.shape[0] tot_images += images.shape[0] + + for label, original_label in zip(labels, original_labels): + assert_that(label).is_equal_to(original_label) + + if SAVE_IMAGES: + save_image(make_grid(ch.concat([images, original_images])/255., images.shape[0]), + os.path.join(IMAGES_TMP_PATH, str(uuid.uuid4()) + '.jpeg') + ) + assert_that(tot_indices).is_equal_to(len(my_dataset)) assert_that(tot_images).is_equal_to(len(my_dataset)) def test_cutout(): - run_test(100, [ - SimpleRGBImageDecoder(), - Cutout(8), - ToTensor() - ], True) + for comp in [True, False]: + run_test(100, [ + SimpleRGBImageDecoder(), + Cutout(8), + ToTensor(), + ToTorchImage() + ], comp) + + +def test_flip(): + for comp in [True, False]: + run_test(100, [ + SimpleRGBImageDecoder(), + RandomHorizontalFlip(1.0), + ToTensor(), + ToTorchImage() + ], comp) + + +def test_module_wrapper(): + for comp in [True, False]: + run_test(100, [ + SimpleRGBImageDecoder(), + ToTensor(), + ToTorchImage(), + ModuleWrapper(tvt.Grayscale(3)), + ], comp) - run_test(100, [ - SimpleRGBImageDecoder(), - Cutout(8), - ToTensor() - ], False) + +def test_mixup(): + for comp in [True, False]: + run_test(100, [ + SimpleRGBImageDecoder(), + ImageMixup(.5, False), + ToTensor(), + ToTorchImage() + ], comp) def test_poison(): @@ -62,11 +118,94 @@ def test_poison(): # Red sqaure mask[:5, :5, 0] = 1 alpha = np.ones((32, 32)) + + for comp in [True, False]: + run_test(100, [ + SimpleRGBImageDecoder(), + Poison(mask, alpha, list(range(100))), + ToTensor(), + ToTorchImage() + ], comp) + + +def test_random_resized_crop(): + for comp in [True, False]: + run_test(100, [ + SimpleRGBImageDecoder(), + RandomResizedCrop(scale=(0.08, 1.0), + ratio=(0.75, 4/3), + size=32), + ToTensor(), + ToTorchImage() + ], comp) + + +def test_translate(): + for comp in [True, False]: + run_test(100, [ + SimpleRGBImageDecoder(), + RandomTranslate(padding=10), + ToTensor(), + ToTorchImage() + ], comp) + + +## Torchvision Transforms +def test_torchvision_greyscale(): + run_test(100, [ + SimpleRGBImageDecoder(), + ToTensor(), + ToTorchImage(), + tvt.Grayscale(3), + ]) + +def test_torchvision_centercrop_pad(): + run_test(100, [ + SimpleRGBImageDecoder(), + ToTensor(), + ToTorchImage(), + tvt.CenterCrop(10), + tvt.Pad(11) + ]) + +def test_torchvision_random_affine(): run_test(100, [ SimpleRGBImageDecoder(), - Poison(mask, alpha, [0, 1, 2]), - ToTensor() - ], False) + ToTensor(), + ToTorchImage(), + tvt.RandomAffine(25), + ]) + +def test_torchvision_random_crop(): + run_test(100, [ + SimpleRGBImageDecoder(), + ToTensor(), + ToTorchImage(), + tvt.Pad(10), + tvt.RandomCrop(size=32), + ]) + +def test_torchvision_color_jitter(): + run_test(100, [ + SimpleRGBImageDecoder(), + ToTensor(), + ToTorchImage(), + tvt.ColorJitter(.5, .5, .5, .5), + ]) + if __name__ == '__main__': - test_poison() + # test_cutout() + test_flip() + # test_module_wrapper() + # test_mixup() + # test_poison() + # test_random_resized_crop() + # test_translate() + + ## Torchvision Transforms + # test_torchvision_greyscale() + # test_torchvision_centercrop_pad() + # test_torchvision_random_affine() + # test_torchvision_random_crop() + # test_torchvision_color_jitter() diff --git a/tests/test_traversal_orders.py b/tests/test_traversal_orders.py index 9b771a95..21c00b6e 100644 --- a/tests/test_traversal_orders.py +++ b/tests/test_traversal_orders.py @@ -29,24 +29,29 @@ def __getitem__(self, index): return (index, np.sin(np.array([index])).view(' 1: init_process_group('nccl', sync_url, rank=rank, world_size=world_size) loader = Loader(fname, 8, num_workers=2, order=order, drop_last=False, - distributed=world_size > 1) + distributed=world_size > 1, indices=indices) result = [] for _ in range(3): content = np.concatenate([x[0].numpy().reshape(-1).copy() for x in loader]) result.append(content) result = np.stack(result) + np.save(path.join(out_folder, f"result-{rank}.npy"), result) -def prep_and_run_test(num_workers, order): +def prep_and_run_test(num_workers, order, with_indices=False): length = 600 + indices = None + if with_indices: + indices = np.random.choice(length, length//2, replace=False) + with TemporaryDirectory() as folder: name = path.join(folder, 'dataset.beton') sync_file = path.join(folder, 'share') @@ -58,7 +63,7 @@ def prep_and_run_test(num_workers, order): writer.from_indexed_dataset(dataset) - args = (num_workers, name, order, sync_file, folder) + args = (num_workers, name, order, sync_file, folder, indices) if num_workers > 1: spawn(process_work, nprocs=num_workers, args=args) else: @@ -71,19 +76,22 @@ def prep_and_run_test(num_workers, order): results = np.concatenate(results, 1) + # For each epoch for i in range(results.shape[0]): - if order == OrderOption.SEQUENTIAL and i < results.shape[0] - 1: - assert_that((results[i] == results[i + 1]).all()).is_true() - if order != OrderOption.SEQUENTIAL and i < results.shape[0] - 1: - assert_that((results[i] == results[i + 1]).all()).is_false() - - epoch_content = Counter(results[i]) - indices_gotten = np.array(sorted(list(epoch_content.keys()))) - assert_that(np.all(np.arange(length) == indices_gotten)).is_true() - assert_that(min(epoch_content.values())).is_equal_to(1) - assert_that(max(epoch_content.values())).is_less_than_or_equal_to(2) - - + if not with_indices: + if order == OrderOption.SEQUENTIAL and i < results.shape[0] - 1: + assert_that((results[i] == results[i + 1]).all()).is_true() + if order != OrderOption.SEQUENTIAL and i < results.shape[0] - 1: + assert_that((results[i] == results[i + 1]).all()).is_false() + + epoch_content = Counter(results[i]) + indices_gotten = np.array(sorted(list(epoch_content.keys()))) + assert_that(np.all(np.arange(length) == indices_gotten)).is_true() + assert_that(min(epoch_content.values())).is_equal_to(1) + assert_that(max(epoch_content.values())).is_less_than_or_equal_to(2) + else: + assert_that(set(results[i])).is_equal_to(set(indices)) + def test_traversal_sequential_1(): prep_and_run_test(1, OrderOption.SEQUENTIAL) @@ -123,3 +131,13 @@ def test_traversal_quasirandom_3(): @pytest.mark.skip() def test_traversal_quasirandom_4(): prep_and_run_test(4, OrderOption.QUASI_RANDOM) + +def test_traversal_sequential_distributed_with_indices(): + prep_and_run_test(2, OrderOption.SEQUENTIAL, True) + +def test_traversal_random_distributed_with_indices(): + prep_and_run_test(2, OrderOption.RANDOM, True) + +@pytest.mark.skip() +def test_traversal_quasi_random_distributed_with_indices(): + prep_and_run_test(2, OrderOption.QUASI_RANDOM, True) \ No newline at end of file