$ nvidia-smi # 显示当前GPU使用情况
$ watch -n 1 -d nvidia-smi # 每隔1s刷新一次,时间参数可以更改
torch.cuda.is_available() 返回True表示有GPU。
torch.cuda.device_count() 可以获得能够使用的GPU数量。
把数据从内存转移到GPU,一般针对张量(我们需要的数据)和模型。 对张量(类型为FloatTensor或者是LongTensor等),一律直接使用方法.to(device)或.cuda()即可。
# 定义device,其中需要注意的是“cuda:0”代表起始的device_id为0,如果直接是“cuda”,同样默认是从0开始。
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 数据
for batch_idx, (img, label) in enumerate(train_loader):
img=img.to(device)
label=label.to(device)
# 网络模型
model = Net()
model.to(device)
请注意,仅调用my_tensor.to(device)会在GPU上返回my_tensor的新副本,而不是重写my_tensor。 您需要将其分配给新的张量,并在GPU上使用该张量。即 img=img.to(device),不能简单的写为img.to(device)。
在多个GPU上执行forward, backward propagations是很自然的。 但是,Pytorch默认仅使用一个GPU。单机多GPUs主要采用的DataParallel函数,而不是DistributedParallel,后者一般用于多主机多GPUs,当然也可用于单机多GPUs。
使用时先将model传入torch.nn.DataParallel函数,然后to(device)即可。官方代码文档如下:nn.DataParallel 教程文档如下:tutorial.
model = Net() # 网络模型
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device)
这时,默认所有存在的显卡都会被使用。如果你只想利用其中一部分,如只使用编号为0、1的两个GPU,那么可以采用以下方式,需要注意:device_ids的起始编号要与之前定义的device中的“cuda:0”相一致,不然会报错。
gpu_ids = [1, 2]
model = nn.DataParallel(model,device_ids= gpu_ids)
其实 DataParallel
这个函数有两个可选参数 device_ids, output_device
, 前者是指在哪些 gpu 上运行, 默认是全部可见 gpu 上, 后者是指最终将 loss 汇总到哪个 gpu 上, 默认是 cuda:0 .
服务器首先将数据和模型拷贝到 cuda:0 上, cuda:0 也就是我们前面定义的起始 device. 由于申明了 DataParallel, pytorch 会在多个 gpu 上并行地将 model apply 到 data 上. 为了达到这个目的, 需要完成如下步骤:
多 gpu 运行的过程导致一些问题:
使用GPU可以提升我们训练的速度,如果使用不当,可能影响使用效率,具体使用时要注意以下几点:
1. 保存和加载整个模型。这种方式再重新加载的时候不需要自定义网络结构,保存时已经把网络结构保存了下来,比较死板不能调整网络结构。
torch.save(model, 'model.pkl') # 其实后缀名是啥都行,无所谓
model = torch.load('model.pkl')
2.仅保存和加载模型参数(推荐使用)。这种方式再重新加载的时候需要自己定义网络,并且其中的参数名称与结构要与保存的模型中的一致(可以是部分网络,比如只使用VGG的前几层),相对灵活,便于对网络进行修改。
torch.save(model.state_dict(), 'params.pkl')
model.load_state_dict(torch.load('params.pkl'))
由于多GPU训练使用了 nn.DataParallel(net, device_ids=gpu_ids) 对网络进行封装,因此在原始网络结构中添加了一层module。网络结构如下:
DataParallel(
(module): Net(
(conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(fc1): Linear(in_features=400, out_features=120, bias=True)
(fc2): Linear(in_features=120, out_features=84, bias=True)
(fc3): Linear(in_features=84, out_features=10, bias=True)
)
)
而不使用多GPU训练的网络结构如下:
Net(
(conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(fc1): Linear(in_features=400, out_features=120, bias=True)
(fc2): Linear(in_features=120, out_features=84, bias=True)
(fc3): Linear(in_features=84, out_features=10, bias=True)
)
由于在测试模型时不需要用到多GPU测试,因此在保存模型时应该把module层去掉。如下:
if len(gpu_ids) > 1:
torch.save(model.module.state_dict(), "model.pth")
# torch.save(model.module,'model.pkl') # 保存整个模型
else:
t.save(model.state_dict(), "model.pth")
cpu加载GPU训练的模型 ,可采用map_location='cpu',如果采用多GPU训练的模型保存时忘记去掉module层,可以加载模型后,提取module层内的内容重新赋给模型。
model = torch.load('xxx.pkl', map_location='cpu')
model = model.module # 去掉module层
使用Remote-SSH插件