PyTorch 乱数固定

PyTorchで深層学習を行ったときの再現性を担保するために乱数固定を行うベストプラクティス．

import os
import random

seed = 0

def seed_everything(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed) # pythonのハッシュベース操作の再現性担保
    np.random.seed(seed)
    torch.manual_seed(seed) # ネットワーク重みの初期値を固定
    torch.cuda.manual_seed(seed) # ネットワーク重みの初期値を固定 (GPU)

    # cuDNN: NVIDIAのConvolution高速化ライブラリ．
    torch.backends.cudnn.deterministic = True # cuDNNによる最適化プロセス固定
    torch.backends.cudnn.benchmark = True # input_size固定CNNなら高速化が期待できる．iterationごとに変わるならFalseの方が速い．

seed_everything(seed=seed)