The question I’m about to ask is probably not PyTorch-specific, but I encountered it in context of PyTorch DataLoader.
How do you properly add random perturbations when data is loaded and augmented by several processes?
Let me show on a simple example that this is not a trivial question. I have two files:
- augmentations.py:
import numpy as np
import os
class RandomAugmentation:
def __call__(self, obj):
perturbation = np.random.randint(10)
print(os.getpid(), perturbation)
return obj
- main.py
import numpy as np
from time import sleep
from torch.utils.data import Dataset, DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, ToTensor, Resize
from augmentations import RandomAugmentation
PATH = ...
transform = Compose([
RandomAugmentation(),
Resize((16, 16)),
ToTensor()
])
ds = ImageFolder(PATH, transform=transform)
dl = DataLoader(ds, batch_size=2, num_workers=3)
for epoch_nr in range(2):
for batch in dl:
break
sleep(1)
print('-' * 80)
In the main.py we generate batches of data like we would in a regular image recognition task. RandomAugmentation prints output of RNG that would be used for data augmentation in a real task.
On my machine the output from running main.py was:
20909 6
20909 8
20908 6
20910 6
20908 8
20910 8
20909 7
20909 6
20908 7
20908 6
20910 7
20908 0
20910 6
20908 1
--------------------------------------------------------------------------------
20952 6
20953 6
20952 8
20952 7
20953 8
20952 6
20954 6
20952 0
20952 1
20953 7
20954 8
20953 6
20954 7
20954 6
--------------------------------------------------------------------------------
Clearly, i-th perturbation added by one worker is exactly the same as i-th perturbation added by any other worker (in this example 6, 8, 7, 6, and so on). What’s even worse, in the next epoch data loader is kind of reset, and generates exactly the same (perturbated) data as before, what completely ruins the entire idea of data augmentation.
For the time being I designed such a workaround:
augmentations.py
import numpy as np
from time import time
import os
RNG_PID = None
RNG = None
def get_rng():
global RNG
global RNG_PID
if os.getpid() != RNG_PID:
RNG = np.random.RandomState([int(time()), os.getpid()])
RNG_PID = os.getpid()
print("Initialize RNG", int(time()), os.getpid())
return RNG
class RandomAugmentation:
def __call__(self, obj):
rng = get_rng()
perturbation = rng.randint(10)
print(os.getpid(), perturbation)
return obj
However, it doesn’t seem like a very elegant solution to me.
It would be great to hear how you deal with this matter 
