pytorch中神經網絡的多線程數設置：torch.set_num_threads(N)

實驗室的同學一直都是在服務器上既用CPU訓練神經網絡也有使用GPU的，最近才發現原來在pytorch中可以通過設置  torch.set_num_threads(args.thread)  來限制CPU上進行深度學習訓練的線程數。

 

 

 

torch.set_num_threads(args.thread)   在使用時的一個注意事項就是如果不設置則默認使用物理CPU核心數的線程進行訓練，而往往默認設置是可以保證運算效率最高的，因此該設置線程數是需要小於物理CPU核心數的，否則會造成效率下降。

 

既然默認設置既可以保證最高的運算效率那么這個設置的意義在哪呢，這個設置的意義就是在多人使用計算資源時限制你個人的改應用的計算資源占用情況，否則很可能你一個進程跑起來開了太多的線程直接把CPU占用率搞到50%或者直接奔100%去了。

 

總的說，該設置是為了在多人共享計算資源的時候防止一個進程搶占過高CPU使用率的。

 

 

 

給一個自己的設置代碼：（實現了pytorch的最大可能性的確定性可復現性，並設置訓練、推理時最大的線程數）

 

# pytorch的運行設備
device = None


def context_config(args):
    global device

    seed = args.seed

    random.seed(seed)
    np.random.seed(seed)

    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    torch.set_num_threads(args.thread)  # 設置pytorch並行線程數
    if torch.cuda.is_available() and args.gpu >= 0:
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True

        device = torch.device('cuda:' + str(args.gpu))
    else:
        device = torch.device('cpu')

 

 

 

 

 

 

 

 

================================

 

 

 

運行DQN2013算法，CPU運行，CPU為Intel 10700k, 8核心16線程：

 

默認設置：

 

 

 

設置20線程：

 torch.set_num_threads(20）

 

 

 

 

 

結果：

發現如果線程設置過多，超過CPU的物理線程數運行效率不僅沒有提升反而下降，正常默認設置即可。

 

 

 

 

 

  ===================================