PyTorch dataloader 中的隨機數

起因

最近用PyTorch訓練一個新的CNN，需要在輸入的訓練圖像上增加簡單的噪聲，最開始是利用NumPy提供的隨機數函數去添加噪聲，基於PyTorch提供的dataloader進行圖像加載。在使用dataloader時，發現當worker數量大於1時，每個worker生成的隨機數的序列是完全一樣的。雖然每次啓動dataloader，NumPy給出的隨機數序列不同，但是每個worker使用的是一模一樣的隨機數列。這是我不需要的行爲。

思考

簡單思考一下，原因可能是當使用一個以上的worker時，dataloader採用python的multiprocessing package，生成多個進程去加載數據，在生成進程後，各個子進程沒有各自生成自己獨立的seed，而是繼承了父進程的seed，從而各自生成的隨機數列是完全一致的。於是進行了簡單的測試

from __future__ import print_function

from multiprocessing import Process
import os
import numpy as np
import torch

def func(name, n=4):
    np.random.seed( int(os.getpid()) )

    for i in range(n):
        print("%s: %d, np.random.randn(1) = %+f, torch.randn(1) = %+f. " % ( name, i, np.random.randn(1), torch.randn(1) ))

if __name__ == "__main__":
    processes = []
    processes.append( Process(target=func, args=("p1", 4)) )
    processes.append( Process(target=func, args=("p2", 4)) )

    print("Start all the processes.")

    for p in processes:
        p.start()
    
    for p in processes:
        p.join()

    print("All processes joined.")

上述代碼在註釋掉func( )中的np.random.seed( ) 函數的調用時，輸出如下：

而使用np.random.seed( ) 函數後，輸出變爲

可以看到NumPy輸出的隨機數的變化。在上面兩個測試中，同樣使用了PyTorch提供的torch.randn( )函數，可以看到torch.rand( ) 的配置和NumPy是獨立的，它總是輸出相同的序列。當然，與NumPy類似，可以使用torch.manual_seed( ) 函數配置PyTorch使用的seed，下圖爲同時配置NumPy和PyTorch的seed後的輸出情況

此時，兩個進程輸出的隨機數序列是不同的。

Dataloader

現在的問題是如何修改我之前的dataloader，經過幾次實測，發現PyTorch已經內部進行了處理，在dataloader的不同worker內，PyTorch配置了不同的seed，worker之間所使用torch.randn( )輸出的隨機數序列是不同的。但是PyTorch並沒有配置NumPy，所以默認情況下，dataloader的workder使用NumPy的random.randn() 函數將得到一樣的序列。

from __future__ import print_function

import numpy as np
import os

import torch
import torch.utils.data as data

class RandomDataFolder(data.Dataset):
    def __init__(self, n=3):
        super(RandomDataFolder, self).__init__()

        self.data = [ i for i in range(n) ]

    def __len__(self):
        return len( self.data )

    def __getitem__(self, idx):

        a = torch.FloatTensor( (self.data[idx], ) )

        print("a[0] = %f, np.random.rand(1) = %f, torch.randn(1) = %+f, os.getpid() = %d. " % (a[0], np.random.rand(1), torch.randn(1), os.getpid()))

        return a

    def show(self):
        for d in self.data:
            print(d)

if __name__ == "__main__":
    print("Test the random functions with dataloader.")

    # Create the dataset.
    dataset = RandomDataFolder(8)

    print("The original data: ")
    dataset.show()
    print("")

    # Create the dataloader.
    dataloader = data.DataLoader( dataset, \
        batch_size=2, shuffle=False, num_workers=2, drop_last=False )
    
    # import ipdb; ipdb.set_trace()

    # Test.
    print("The actual loaded data.")
    for batchIdx, (data) in enumerate( dataloader ):
        for i in range(data.size()[0]):
            print("batchIdx = %d, data = %f. " % ( batchIdx, data[i, 0] ))
    

輸出如下圖

結論

在PyTorch的dataloader的worker之間，torch.randn( )將輸出不同的隨機序列，NumPy的random.randn( )將輸出相同的隨機序列。