處理pytorch代碼時如何利用所有GPU？

Question

我有2個GPU，使用pytorch代碼時，僅使用一個GPU。 我嘗試了CUDA_VISIBLE_DEVICES=0,1 python xxx.py ，但是發生了

'CUDA_VISIBLE_DEVICES：找不到命令'

問題。 我還嘗試在對象py文件中添加以下行：

 import os    
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"    
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" 

但仍然只使用一個GPU。

Answer 1

您需要分別將訓練數據並行化到每個GPU。 數據並行使用torch.nn.DataParallel實現。 pytorch文檔中的一個示例：

import torch
import torch.nn as nn


class DataParallelModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.block1 = nn.Linear(10, 20)

        # wrap block2 in DataParallel
        self.block2 = nn.Linear(20, 20)
        self.block2 = nn.DataParallel(self.block2)

        self.block3 = nn.Linear(20, 20)

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        return x