单机多GPU训练模型入门指南(torch.nn.DataParallel)

目录

模型部分

1. 指定使用的GPU

2. 使用Torch的数据并行库(将模型搬到GPU上)

3. 保存模型

数据部分

1. 选择GPU

2. 将数据搬到GPU上

3. loss的反向传播修改

4. 如果需要保存loss数据

查看效果

---

本文将介绍模型和数据两部分的处理。

模型部分

1. 指定使用的GPU

1.1 导入os库

import os

 1.2 给服务器上的GPU编号

最好一次性都编上，从0~n-1，n是服务器上的显卡的数量，一般为4or8张卡。

选择一：写在python代码中

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'

选择二：在命令行设置

CUDA_VISIBLE_DEVICES=0,1,2,3

查看有没有设置好，使用命令行命令

echo $CUDA_VISIBLE_DEVICES

1.3 指定自己要使用哪几张卡

device_ids = [0, 1, 2, 3]

注 这个device_ids的列表要在后面作为参数传入

2. 使用Torch的数据并行库(将模型搬到GPU上)

这里要有两行代码

model = torch.nn.DataParallel(model, device_ids=device_ids)
model = model.cuda()

当然，也可以写在一起

model = torch.nn.DataParallel(model, device_ids=device_ids).cuda()

3. 保存模型

原来

model.save_pretrained(path.my_checkpoint)

现在

model.module.save_pretrained(path.my_checkpoint)

注：这里使用了huggingface的transformers库，使用torch.save应该是类似的。

数据部分

注：这里的数据都是训练函数里面的，测试函数里面的应该类似。

1. 选择GPU

device = torch.device("cuda:0")

2. 将数据搬到GPU上

data = data.to(device)

上面那句和下面这句好像是等价的

data = data.cuda()

3. loss的反向传播修改

这个地方也是我遇到的报错最多的地方。

原来

loss.backward()

改法1

loss.sum().backward()

改法2

loss.backward(torch.ones(loss.shape).to(device))

4. 如果需要保存loss数据

所有的loss都要改成

loss.mean().item()

查看效果

使用命令(每0.1秒刷新一次)

watch -n 0.1 nvidia-smi

 蓝色箭头所指是我的进程。

以上。