在多卡上训练的过程为先将模型和数据加载到第一张卡上,然后copy至其他卡。batchsize最好设为卡的整数倍,比如两张卡,bs为2, 那么每张卡分别计算bs=1的结果,在model.forward之后将不同卡上返回的结果合并传回再做下一步计算。
下面是踩坑后总结的需要注意的一些点。
1. 关于cuda与to(device)
- model 与 data 全部采用 to(device)方法来迁移至显存中,其中
model.to(device)
中的device, 默认是cuda:0
就可以。 但是在网络计算过程中的中间数据,要放到哪个设备上是要根据当时在哪一块卡上计算决定的,因此,应将device设置为当前相关数据在的卡,x.to(related_data.device)
。
2. 关于model与model.module
- 在调用model内部定义的变量或函数时,由于已经使用
model = nn.DataParallel(model)
包裹,因此必须要写做model.module
才能正常调用。 但是,调用model.forward()
时必须要直接用model本身,否则无法使用多卡。
3. 关于多卡计算结果合并
- 必须确保forward最后return的均为tensor, 不能使用其他数据结构包裹,也不能返回标量,多个tensor分别返回。
4. 关于nn.ParameterList()不能用于多卡的问题
- 多卡不支持复制属性为List的参数,代码中定义的
nn.ParameterList()
会在多卡计算中无法复制到其他卡上,导致计算时参数为空的问题。需要写成module.register_parameter()
的形式才能随module一起成功复制。示例如下
for k in range(len(self.filters) + 1):
# Weights
H_init = np.log(np.expm1(1 / scale / filters[k + 1]))
H_k = nn.Parameter(torch.ones((n_channels, filters[k + 1], filters[k]))) # apply softmax for non-negativity
torch.nn.init.constant_(H_k, H_init)
self.register_parameter('H_{}'.format(k), H_k)
# Scale factors
a_k = nn.Parameter(torch.zeros((n_channels, filters[k + 1], 1)))
self.register_parameter('a_{}'.format(k), a_k)
# Biases
b_k = nn.Parameter(torch.zeros((n_channels, filters[k + 1], 1)))
torch.nn.init.uniform_(b_k, -0.5, 0.5)
self.register_parameter('b_{}'.format(k), b_k)
5. 其他
- 关于loss的计算部分不能放在model里,不然会产生不在同一张卡无法计算的问题
- 返回的tensor如果在原有网络中是按batch做了某种运算才返回的,这部分也需要加在model外另外做,因为多卡计算返回的是单独计算的结果。
更改为DistributedDataParallel
注意几个点:
- 模型和数据在两张卡上是一模一样的,相当于分别在两张卡上启动了两个python程序
python -m torch.distributed.launch --nproc_per_node=2 脚本 其他命令行参数
- 一定要确保当前的模型和数据都设置了CUDA,在包裹为module之后一定还要加一句
model.cuda()
。全局可以写一个或者torch.cuda.set_device(local_rank)
self.model=nn.parallel.DistributedDataParallel(
self.model,
device_ids=[opt.local_rank],
output_device=opt.local_rank,
broadcast_buffers=False,
)
网友评论