add option to seed numba rng

Author: juliustao

ffcv/transforms/cutout.py

@@ -21,18 +21,26 @@ class Cutout(Operation):
         Useful for when a normalization layer follows cutout, in which case
         you can set the fill such that the square is zero
         post-normalization.
+    seed : int, optional
+        Random seed set inside code passed to numba (for reproducibility).
+        Nonnegative values are reduced modulo 2**32 to be in valid range.
+        Negative values are not used. Defaults to -1.
     """
-    def __init__(self, crop_size: int, fill: Tuple[int, int, int] = (0, 0, 0)):
+    def __init__(self, crop_size: int, fill: Tuple[int, int, int] = (0, 0, 0), seed: int = -1):
         super().__init__()
         self.crop_size = crop_size
         self.fill = np.array(fill)
+        self.seed = seed % 2 ** 32 if seed >= 0 else seed
 
     def generate_code(self) -> Callable:
         my_range = Compiler.get_iterator()
         crop_size = self.crop_size
         fill = self.fill
+        seed = self.seed
 
         def cutout_square(images, *_):
+            if seed >= 0:
+                np.random.seed(seed)
             for i in my_range(images.shape[0]):
                 # Generate random origin
                 coord = (

ffcv/transforms/flip.py

@@ -1,6 +1,7 @@
 """
 Random horizontal flip
 """
+import numpy as np
 from numpy import dtype
 from numpy.random import rand
 from typing import Callable, Optional, Tuple
@@ -18,17 +19,25 @@ class RandomHorizontalFlip(Operation):
     flip_prob : float
         The probability with which to flip each image in the batch
         horizontally.
+    seed : int, optional
+        Random seed set inside code passed to numba (for reproducibility).
+        Nonnegative values are reduced modulo 2**32 to be in valid range.
+        Negative values are not used. Defaults to -1.
     """
 
-    def __init__(self, flip_prob: float = 0.5):
+    def __init__(self, flip_prob: float = 0.5, seed: int = -1):
         super().__init__()
         self.flip_prob = flip_prob
+        self.seed = seed % 2 ** 32 if seed >= 0 else seed
 
     def generate_code(self) -> Callable:
         my_range = Compiler.get_iterator()
         flip_prob = self.flip_prob
+        seed = self.seed
 
         def flip(images, dst):
+            if seed >= 0:
+                np.random.seed(seed)
             should_flip = rand(images.shape[0]) < flip_prob
             for i in my_range(images.shape[0]):
                 if should_flip[i]:

ffcv/transforms/random_resized_crop.py

@@ -20,16 +20,24 @@ class RandomResizedCrop(Operation):
         Lower and upper bounds for random aspect ratio of the crop.
     size : int
         Side length of the output.
+    seed : int, optional
+        Random seed set inside code passed to numba (for reproducibility).
+        Nonnegative values are reduced modulo 2**32 to be in valid range.
+        Negative values are not used. Defaults to -1.
     """
-    def __init__(self, scale: Tuple[float, float], ratio: Tuple[float, float], size: int):
+    def __init__(self, scale: Tuple[float, float], ratio: Tuple[float, float], size: int, seed: int = -1):
         super().__init__()
         self.scale = scale
         self.ratio = ratio
         self.size = size
+        self.seed = seed % 2 ** 32 if seed >= 0 else seed
 
     def generate_code(self) -> Callable:
         scale, ratio = self.scale, self.ratio
+        seed = self.seed
         def random_resized_crop(im, dst):
+            if seed >= 0:
+                np.random.seed(seed)
             i, j, h, w = fast_crop.get_random_crop(im.shape[0],
                                                 im.shape[1],
                                                 scale,

ffcv/transforms/translate.py

@@ -20,18 +20,26 @@ class RandomTranslate(Operation):
         Max number of pixels to translate in any direction.
     fill : tuple
         An RGB color ((0, 0, 0) by default) to fill the area outside the shifted image.
+    seed : int, optional
+        Random seed set inside code passed to numba (for reproducibility).
+        Nonnegative values are reduced modulo 2**32 to be in valid range.
+        Negative values are not used. Defaults to -1.
     """
 
-    def __init__(self, padding: int, fill: Tuple[int, int, int] = (0, 0, 0)):
+    def __init__(self, padding: int, fill: Tuple[int, int, int] = (0, 0, 0), seed: int = -1):
         super().__init__()
         self.padding = padding
         self.fill = np.array(fill)
+        self.seed = seed % 2 ** 32 if seed >= 0 else seed
 
     def generate_code(self) -> Callable:
         my_range = Compiler.get_iterator()
         pad = self.padding
+        seed = self.seed
 
         def translate(images, dst):
+            if seed >= 0:
+                np.random.seed(seed)
             n, h, w, _ = images.shape
             # y_coords = randint(low=0, high=2 * pad + 1, size=(n,))
             # x_coords = randint(low=0, high=2 * pad + 1, size=(n,))