pytorch-多GPU訓練（單機多卡、多機多卡）

pytorch 單機多卡訓練

首先是數據集的分布處理

需要用到的包：
torch.utils.data.distributed.DistributedSampler
torch.utils.data.DataLoader
torch.utils.data.Dataset

DistributedSampler這個包我們用來確保dataloader只會load到整個數據集的一個特定子集，為每一個子進程劃分出一部分數據集，以避免不同進程之間數據重復。

而我們構造的Sampler是在構造DataLoader時的一個參數

先構建一個自己的torch.utils.data.Dataset類的dataset類:

import torch.utils.data as Data class MyDataSet(Data.Dataset):def __init__(self, x_train, x_label, y_train, y_label): super(MyDataSet, self).__init__()self.x_train = x_trainself.x_label = x_labelself.y_train = y_trainself.y_label = y_labeldef __len__(self):return self.x_train.shape[0]def __getitem__(self, idx):return self.x_train[idx], self.x_label[idx], self.y_train[idx], self.y_label[idx]

初始化DataSet:
這里的x_train, x_label, y_train, y_label是數據集加載進來的變量，根據自己的數據集情況設置變量和數量

train_data = MyDataSet(x_train, x_label, y_train, y_label)

然后根據構造的Dataset對象構造一個DistributedSampler：

train_sampler = torch.utils.data.distributed.DistributedSampler(train_data, shuffle=True)

然后使用torch.utils.data.DataLoader構造自己的DataLoader

MyDataLoader = Data.DataLoader(MyDataSet(x_train, x_label, \y_train, y_label), batch_size=2, sampler = train_sampler)

然后是模型的分布處理

需要用到的包：
torch.nn.parallel.DistributedDataParallel
在初始化model之后，進行一步：

model = torch.nn.parallel.DistributedDataParallel(model)# device_ids will include all GPU devices by default

pytorch 多機多卡訓練

多機多卡需要配置多機間通訊，進行多進程的初始化：
需要用到的包：
torch.distributed
在初始化model之后，進行一步：

torch.distributed.init_process_group(init_method='tcp://{}:{}'.format(args.ip, args.port),backend="gloo",

參數解析：

backend: 多個機器之間交換數據的協議 init_method: 機器之間交換數據, 需要指定一個主節點，這代表主節點 world_size: 進程個數,實際上就是機器的個數, 比如四臺臺機器一起訓練的話, world_size就為4 rank: 區分主節點和從節點的, 主節點為0, 剩余的為了1-(N-1), N為要使用的機器的數量, 也就是world_size 創作挑戰賽新人創作獎勵來咯，堅持創作打卡瓜分現金大獎

總結

以上是生活随笔為你收集整理的pytorch-多GPU训练（单机多卡、多机多卡）的全部內容，希望文章能夠幫你解決所遇到的問題。

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