pytorc使用多个GPU同时训练模型

pytorch使用同一设备上多个GPU同时训练模型，只需在原有代码中稍作修改即可。

改动1：

1 2	`<code-pre` `class="code-pre code-pre-line"` `id="pre-GeKjJe"><code-line` `class="line-numbers-rows"></code-line>os.environ['CUDA_VISIBLE_DEVICES']` `=` `'0,1,3'` `# 这里输入你的GPU_id` `</code-pre>`

改动2：

<code-pre class="code-pre code-pre-line" id="pre-6EdzGk"><code-line class="line-numbers-rows"></code-line>if torch.cuda.device_count() > 1:
<code-line class="line-numbers-rows"></code-line>  model = nn.DataParallel(model)
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>model.to(device)
</code-pre>

使用多GPU训练，速度明显得到提升。

官方示例代码

<code-pre class="code-pre code-pre-line" id="pre-frCRCM"><code-line class="line-numbers-rows"></code-line>import torch
<code-line class="line-numbers-rows"></code-line>import torch.nn as nn
<code-line class="line-numbers-rows"></code-line>from torch.utils.data import Dataset, DataLoader
<code-line class="line-numbers-rows"></code-line>import os 
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>os.environ['CUDA_VISIBLE_DEVICES'] = '2,3' # 这里输入你的GPU_id
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line># Parameters and DataLoaders
<code-line class="line-numbers-rows"></code-line>input_size = 5
<code-line class="line-numbers-rows"></code-line>output_size = 2
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>batch_size = 30
<code-line class="line-numbers-rows"></code-line>data_size = 100
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line># Dummy DataSet
<code-line class="line-numbers-rows"></code-line>class RandomDataset(Dataset):
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>    def __init__(self, size, length):
<code-line class="line-numbers-rows"></code-line>        self.len = length
<code-line class="line-numbers-rows"></code-line>        self.data = torch.randn(length, size)
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>    def __getitem__(self, index):
<code-line class="line-numbers-rows"></code-line>        return self.data[index]
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>    def __len__(self):
<code-line class="line-numbers-rows"></code-line>        return self.len
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
<code-line class="line-numbers-rows"></code-line>                         batch_size=batch_size, shuffle=True)
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line># Simple Model
<code-line class="line-numbers-rows"></code-line>class Model(nn.Module):
<code-line class="line-numbers-rows"></code-line>    # Our model
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>    def __init__(self, input_size, output_size):
<code-line class="line-numbers-rows"></code-line>        super(Model, self).__init__()
<code-line class="line-numbers-rows"></code-line>        self.fc = nn.Linear(input_size, output_size)
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>    def forward(self, input):
<code-line class="line-numbers-rows"></code-line>        output = self.fc(input)
<code-line class="line-numbers-rows"></code-line>        print("\tIn Model: input size", input.size(),
<code-line class="line-numbers-rows"></code-line>              "output size", output.size())
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>        return output
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line># Create Model and DataParallel
<code-line class="line-numbers-rows"></code-line>model = Model(input_size, output_size)
<code-line class="line-numbers-rows"></code-line>if torch.cuda.device_count() > 1:
<code-line class="line-numbers-rows"></code-line>  print("Let's use", torch.cuda.device_count(), "GPUs!")
<code-line class="line-numbers-rows"></code-line>  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
<code-line class="line-numbers-rows"></code-line>  model = nn.DataParallel(model)
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>model.to(device)
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>
<code-line class="line-numbers-rows"></code-line>#Run the Model
<code-line class="line-numbers-rows"></code-line>for data in rand_loader:
<code-line class="line-numbers-rows"></code-line>    input = data.to(device)
<code-line class="line-numbers-rows"></code-line>    output = model(input)
<code-line class="line-numbers-rows"></code-line>    print("Outside: input size", input.size(),
<code-line class="line-numbers-rows"></code-line>          "output_size", output.size())
</code-pre>

__EOF__

本文作者：lishuaics
本文链接：https://www.cnblogs.com/L-shuai/p/15811569.html
关于博主：IT小白
版权声明：本博客所有文章除特别声明外，均采用 BY-NC-SA 许可协议。转载请注明出处！
声援博主：如果您觉得文章对您有帮助，可以点击文章右下角【推荐】一下。您的鼓励是博主的最大动力！