Pytorch下使用多GPU訓練模型

注：本文針對單個伺服器上多塊GPU的使用，不是多伺服器多GPU的使用。

當你的電腦中有多塊GPU時，可以增加batch_size來加快訓練速度等。

涉及的代碼

import torch

1. 設定可以使用的GPU編号

   import os

   os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'

   # 使用第0和第3編号的GPU
2. 判斷你的電腦或者pytorch是否能使用GPU來加速模型的訓練

   torch.cuda.is_available()

   -->

   model = model.cuda()

   或者

   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

   -->

   model = model.to(device)

3. 判斷你的電腦有幾塊GPU

   torch.cuda.device_count()

4. torch.nn.DataParallel啟用多GPU并行計算

   torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

   說明：

   該函數實作了在module級别上的資料并行使用，注意batch size要大于GPU的數量。

   參數 ：

   module：需要多GPU訓練的模型

   device_ids&