- 1 使用单块GPU训练模型
- 2 使用多卡(多GPU)训练模型
- 2.1 nn.DataParallel原理
- 2.2 实现代码(亲测有效)
- 2.3 贴图展示
- 2.4 模型保存和加载
1 使用单块GPU训练模型
下面代码主要实现： 如何将tensor放入指定GPU上计算、如何将py文件放入指定GPU上执行
import os
import torch
from torchvision import models
a = torch.randn(3, 3)
b = models.vgg16()
# 1.判断是否有GPU资源
print("是否有GPU：", torch.cuda.is_available())
print("GPU设备数量: ", torch.cuda.device_count())
# 2.指定局部部分tensor使用GPU
if torch.cuda.is_available():
a.cuda() # 法1
b.cuda(1)
device = torch.device("cuda: 1") # 法2
c = torch.randn(3, 3, device = device, requires_grad = True)
# 3.指定全局使用哪一块GPU
# 法1-在终端执行脚本时直接指定GPU的方式
CUDA_VISIBLE_DEVICES=2 python3 train.py
# 法2-脚本前面放入以下代码
torch.cuda.set_device(1)
# 法3-脚本前面放入以下代码
os.environ["CUDA_VISIBLE_DEVICES"] = '1'2 使用多卡(多GPU)训练模型
这里主要讨论单机多卡，单机多卡借助pytorch的nn.DataParallel，多机多卡借助pytorch的DistributedSampler和DistributedDataParallel
对于多机多卡暂不讨论！
2.1 nn.DataParallel原理
首先在前向过程中，你的输入数据会被划分成多个子部分（以下称为副本）送到不同的device中进行计算，而你的模型module是在每个device上进行复制一份，也就是说，输入的batch是会被平均分到每个device中去，但是你的模型module是要拷贝到每个devide中去的，每个模型module只需要处理每个副本即可，当然你要保证你的batch size大于你的gpu个数。然后在反向传播过程中，每个副本的梯度被累加到原始模块中。概括来说就是：DataParallel 会自动帮我们将数据切分并 load 到相应 GPU，将模型复制到相应 GPU，每张卡进行正向传播计算梯度并汇总到主卡上，主卡更新好模型权重以后把权重分发到其余卡。
pytorch实现单机多卡十分容易，其基本原理就是：加入我们一次性读入一个batch的数据, 其大小为[16, 10, 5]，我们有四张卡可以使用。那么计算过程遵循以下步骤：
1.假设我们有4个GPU可以用，pytorch先把模型同步放到4个GPU中。
2.那么首先将数据分为4份，按照次序放置到四个GPU的模型中，每一份大小为[4, 10, 5]；
3.每个GPU分别进行前项计算过程；
4.前向过程计算完后，pytorch再从四个GPU中收集计算后的结果假设[4, 10, 5]，然后再按照次序将其拼接起来[16, 10, 5]，计算loss。
整个过程其实就是 同步模型参数→分别前向计算→计算损失→梯度反传
2.2 实现代码(亲测有效)
指定多卡训练模型
# 假设有四张卡
os.environ["CUDA_VISIBLE_DEVICES"] = '0, 2, 1, 3' # 程序开头必须加上这行代码，这行代码是告诉当前py文件有四块GPU可用
mymodel = ...
model = mymodel()
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model).cuda()2.3 贴图展示
2.4 模型保存和加载
指定多卡训练的模型就不是原来模型的类型了，而是并行化后的模型:
可以打印出来看
并行化后的模型参数必须加载到并行化的模型中，没并行化的参数要加载到没并行化的模型中，不然会出bug。
那么并行化后，如何保存没并行化的结果呢(代码如下)
torch.save(model.module.state_dict(), 'model_name.pth').
