V0.0.3

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']

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()

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 = {

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',

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

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))

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))

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')))

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)

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))
+

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))]

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))]

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',