1. 程式人生 > >PyTorch的可重複性問題 (如何使實驗結果可復現)

PyTorch的可重複性問題 (如何使實驗結果可復現)

由於在模型訓練的過程中存在大量的隨機操作,使得對於同一份程式碼,重複執行後得到的結果不一致。因此,為了得到可重複的實驗結果,我們需要對隨機數生成器設定一個固定的種子。

許多部落格都有介紹如何解決這個問題,但是很多都不夠全面,往往不能保證結果精確一致。我經過許多調研和實驗,總結了以下方法,記錄下來。

全部設定可以分為三部分:

1. CUDNN

cudnn中對卷積操作進行了優化,犧牲了精度來換取計算效率。如果需要保證可重複性,可以使用如下設定:

from torch.backends import cudnn
cudnn.benchmark = False            # if benchmark=True, deterministic will be False
cudnn.deterministic = True

不過實際上這個設定對精度影響不大,僅僅是小數點後幾位的差別。所以如果不是對精度要求極高,其實不太建議修改,因為會使計算效率降低。

2. Pytorch

torch.manual_seed(seed)            # 為CPU設定隨機種子
torch.cuda.manual_seed(seed)       # 為當前GPU設定隨機種子
torch.cuda.manual_seed_all(seed)   # 為所有GPU設定隨機種子

3. Python & Numpy

如果讀取資料的過程採用了隨機預處理(如RandomCrop、RandomHorizontalFlip等),那麼對python、numpy的隨機數生成器也需要設定種子。

import random
import numpy as np
random.seed(seed)
np.random.seed(seed)

最後,關於dataloader:

注意,如果dataloader採用了多執行緒(num_workers > 1), 那麼由於讀取資料的順序不同,最終執行結果也會有差異。也就是說,改變num_workers引數,也會對實驗結果產生影響。目前暫時沒有發現解決這個問題的方法,但是隻要固定num_workers數目(執行緒數)不變,基本上也能夠重複實驗結果。

對於不同執行緒的隨機數種子設定,主要通過DataLoader的worker_init_fn引數來實現。預設情況下使用執行緒ID作為隨機數種子。如果需要自己設定,可以參考以下程式碼:

GLOBAL_SEED = 1

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

GLOBAL_WORKER_ID = None
def worker_init_fn(worker_id):
    global GLOBAL_WORKER_ID
    GLOBAL_WORKER_ID = worker_id
    set_seed(GLOBAL_SEED + worker_id)

dataloader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=2, worker_init_fn=worker_init_fn)