python - 如何在 pytorch 中使用多个 GPU 训练模型？

Question

我的服务器有两个 GPU，如何同时使用两个 GPU 进行训练以最大化它们的计算能力？我下面的代码正确吗？它是否允许我的模型得到适当的训练？

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.bert = pretrained_model
        # for param in self.bert.parameters():
        #     param.requires_grad = True
        self.linear = nn.Linear(2048, 4)


    #def forward(self, input_ids, token_type_ids, attention_mask):
    def forward(self, input_ids, attention_mask):
        batch = input_ids.size(0)
        #output = self.bert(input_ids, token_type_ids, attention_mask).pooler_output
        output = self.bert(input_ids, attention_mask).last_hidden_state
        print('last_hidden_state',output.shape) # torch.Size([1, 768]) 
        #output = output.view(batch, -1) #
        output = output[:,-1,:]#(batch_size, hidden_size*2)(batch_size,1024)
        output = self.linear(output)
        return output

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
    print("Use", torch.cuda.device_count(), 'gpus')
    model = MyModel()
    model = nn.DataParallel(model)
    model = model.to(device)

Answer 1

在多个 GPU 上训练有两种不同的方法:

1. 数据并行性 = 将无法放入单个 GPU 内存的大批量拆分为多个 GPU，因此每个 GPU 将处理可放入其 GPU 的小批量
2. 模型并行性 = 将模型中的层拆分到不同的设备中有点难以管理和处理。

Please refer to this post for more information

在纯 PyTorch 中做数据并行，请引用 this example我创建了一段时间的 PyTorch 的最新更改(截至今天，1.12)。

要利用其他库进行多 GPU 训练而无需设计很多东西，我建议使用 PyTorch Lightning因为它有一个简单的 API 和良好的文档来学习如何使用 Data Parallelism 进行多 GPU 训练。

---

更新:2022/10/25

这是一个视频，详细解释了不同类型的分布式训练:https://youtu.be/BPYOsDCZbno?t=1011