Search code examples
pythondeep-learningneural-networkpytorchtorchvision

Fix random seed for torchvision transforms


I use some code similar to the following - for data augmentation:

    from torchvision import transforms

    #...

    augmentation = transforms.Compose([
        transforms.RandomApply([
            transforms.RandomRotation([-30, 30])
        ], p=0.5),
        transforms.RandomHorizontalFlip(p=0.5),
    ])

During my testing I want to fix random values to reproduce the same random parameters each time I change the model training settings. How can I do it?

I want to do something similar to np.random.seed(0) so each time I call random function with probability for the first time, it will run with the same rotation angle and probability. In other words, if I do not change the code at all, it must reproduce the same result when I rerun it.

Alternatively I can separate transforms, use p=1, fix the angle min and max to a particular value and use numpy random numbers to generate results, but my question if I can do it keeping the code above unchanged.


Solution

  • In the __getitem__ of your dataset class make a numpy random seed.

    def __getitem__(self, index):      
        img = io.imread(self.labels.iloc[index,0])
        target = self.labels.iloc[index,1]
    
        seed = np.random.randint(2147483647) # make a seed with numpy generator 
        random.seed(seed) # apply this seed to img transforms
        if self.transform is not None:
            img = self.transform(img)
    
        random.seed(seed) # apply this seed to target transforms
        if self.target_transform is not None:
            target = self.target_transform(target)
    
        return img, target