Pytorch多GPU運行

設置可用GPU環(huán)境變量。例如，使用0號和1號GPU'

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'

設置模型參數(shù)放置到多個GPU上。在pytorch1.0之后的版本中，多GPU運行變得十分方便，先將模型的參數(shù)設置并行

    if torch.cuda.device_count() > 1:
      print("Let's use", torch.cuda.device_count(), "GPUs!")
      model = nn.DataParallel(model)

將模型參數(shù)設置使用GPU運行

    if torch.cuda.is_available():
      model.cuda()

踩坑記錄

在訓練中，需要使用驗證集/測試集對目前的準確率進行測試，驗證集/測試集的加載也會占用部分顯存，所以在訓練開始時，不要將所有顯存都幾乎占滿，稍微留一些顯存給訓練過程中的測試環(huán)節(jié)

pytorch并行后，假設batchsize設置為64，表示每張并行使用的GPU都使用batchsize=64來計算（單張卡使用時，使用batchsize=64比較合適時，多張卡并行時，batchsize仍為64比較合適，而不是64*并行卡數(shù)）。

參考

https://www.zhihu.com/question/67726969

以上就是本文的全部內(nèi)容，希望對大家的學習有所幫助，也希望大家多多支持腳本之家。