如何使用GPU訓練/測試模型

• 使用單GPU
• • 設置設備
  • 將數據轉換成CUDA張量
  • 將模型參數轉換成CUDA張量
  • 使用指定GPU
  • • 1.使用CUDA_VISIBLE_DEVICES。
    • • • 1.1 直接在終端或shell腳本中設定：
        • 1.2 python代碼中設定：
    • 2. 使用函數 set_device
• 使用多GPU
• • DP方法
  • DDP方法
  • • 需要先初始化
    • 數據集的處理
    • 模型初始化
    • 單節點多GPU分布式訓練
• 實驗結果

原理：通過依靠GPU的并行計算能力，能夠大大縮短模型訓練時間。
在使用GPU跑代碼的時候，只需要將模型參數和數據放到GPU上轉換成CUDA張量即可。

所以，代碼需要修改的地方只有兩處：

• 1.模型實例化處。
• 2.數據迭代處(如果數據有更新或增加的話也需要在相應的地方改，只要確定在輸入模型前的所有數據都轉換為CUDA張量即可)。

使用單GPU

設置設備

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

cuda:0表示使用0號顯卡進行訓練，所以如果需要指定其他GPU，則可以直接修改該數字即可。如cuda:1，表示使用1號顯卡進行訓練。

將數據轉換成CUDA張量

也就是將數據放到GPU上

for i,data in enumerate(trainloader,0):#從迭代器中獲取數據輸入inputs,labels = datainputs, labels = inputs.to(device), labels.to(device) #數據轉換為CUDA張量outputs=net(inputs) #數據輸入模型

to(device)的作用是將數據轉換為CUDA張量。

以一個labels為例：

• 加.to(device)之前，labels為tensor([0, 6, 2, 0])
• 加.to(device)之后，labels為tensor([0, 6, 2, 0], device=‘cuda:0’)

將模型參數轉換成CUDA張量

同理，也就是將模型放到GPU上

net=net.to(device) #model為實例化的模型

使用指定GPU

PyTorch默認使用從0開始的GPU，如果GPU0正在運行程序，需要指定其他GPU。

除了上面直接修改cuda:0的方法之外還有以下兩種方法來指定需要使用的GPU。

1.使用CUDA_VISIBLE_DEVICES。

1.1 直接在終端或shell腳本中設定：

CUDA_VISIBLE_DEVICES=1 python my_script.py

1.2 python代碼中設定：

import os os.environ["CUDA_VISIBLE_DEVICES"] = "2"

2. 使用函數 set_device

import torch torch.cuda.set_device(id)

該函數見 pytorch-master\torch\cuda_init_.py。
不過官方建議使用CUDA_VISIBLE_DEVICES，不建議使用 set_device 函數。

使用多GPU

DP方法

如果有多個GPU，使用nn.DataParallel來包裝我們的模型。 然后通過model.to(device)把模型放到GPU上。
代碼如下：

device=torch.device("cuda" if torch.cuda.is_available() else "cpu") net=Net().to(device) #需要先加載到GPU上，將模型參數轉換成CUDA向量 if torch.cuda.device_count() > 1:net=nn.DataParallel(net)

DDP方法

DDP方法比DP方法要好，解決了DP數據分配不平衡和訓練速度慢的缺點。
需要添加的地方有：

需要先初始化

需初始化local_rank參數，這里通過argparse模塊設置：

parser.add_argument('--local_rank', type=int, default=0,help='node rank for distributed training')

local_rank參數代表要訓練的機子，本來用于標記主機和從機的，如果是多機的話，不同的機器使用不同的local_rank標識機器，由于這里是單機多卡，只用0表示主機就可以了。
另外還需要初始化進程組，代碼如下：

torch.distributed.init_process_group(backend="nccl", init_method='file://file:///DATA/wanghongzhi/first_dnn/temp/test',rank=0, world_size=1)

其中rank是主機的編號，world_size表示分布式主機的個數。

數據集的處理

這步主要是保證一個batch里的數據被均攤到進程上，每個進程能獲取到相應的數據。每個GPU產生一個進程。

train_sampler = torch.utils.data.distributed.DistributedSampler(trainset) trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,shuffle=False,sampler=train_sampler)

模型初始化

先執行net=Net().to(device)將模型放在GPU上后，再執行下面語句初始化模型：

net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank],output_device=args.local_rank)

單節點多GPU分布式訓練

運行代碼的shell腳本為：

python -m torch.distributed.launch --nproc_per_node=2 main.py

其中nproc_per_node為使用的顯卡數量。

關于DDP的更多細節、參數的選擇及其作用可以參考：
博客1
博客2
博客3
博客4
github資源

實驗結果

顯卡使用TITIAN Xp，從下面數據結果看到DDP效果最好。

• 1.使用CPU耗時888.595s
• 2.使用一塊GPU耗時80.06s
• 3.使用DP,兩塊GPU耗時60.7744s
• 4.使用DDP,兩塊GPU耗時48.7986s

參考：
單GPU
多GPU

總結

以上是生活随笔為你收集整理的pytorch使用GPU炼丹笔记的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。