如何用多个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. 模型并行=将模型中的层划分成不同的设备是有点棘手的管理和处理。

有关详情，请参阅本帖。

要在纯PyTorch中实现数据并行，请参考我创建的这个例子，该这个例子回到了PyTorch的最新更改(截至今天，1.12)。

为了利用其他库来进行多GPU培训而不需要进行许多工程，我建议使用PyTorch闪电，因为它有一个简单易懂的API和良好的文档来学习如何使用数据并行性进行多GPU培训。

更新日期: 2022/10/25

下面是一个视频，详细介绍了不同类型的分布式培训：https://youtu.be/BPYOsDCZbno?t=1011

Answer 2

我使用数据并行。我指的是此链接。这是一个有用的参考资料。