A memory balanced and communication efficient FullyConnected layer with CrossEntropyLoss model parallel implementation in PyTorch
在人脸和re-id领域,部分私有的数据集的label数量可达上百万/千万/亿的规模,此时fc层的参数量就足以把显卡容量撑满,导致只能使用较小的batch_size
,训练速度较慢,效果不佳
class FullyConnected(nn.Module):
def __init__(self, in_dim, out_dim, num_gpu, model_parallel=False):
super(FullyConnected, self).__init__()
self.num_gpu = num_gpu
self.model_parallel = model_parallel
if model_parallel:
self.fc_chunks = nn.ModuleList()
for i in range(num_gpu):
_class_num = out_dim // num_gpu
if i < (out_dim % num_gpu):
_class_num += 1
self.fc_chunks.append(
nn.Linear(in_dim, _class_num, bias=False).cuda(i)
)
else:
self.classifier = nn.Linear(in_dim, out_dim, bias=False)
def forward(self, x):
if self.model_parallel:
x_list = []
for i in range(self.num_gpu):
_x = self.fc_chunks[i](x.cuda(i))
x_list.append(_x)
x = torch.cat(x_list, dim=1)
return x
else:
return self.classifier(x)
类似的实现在这个基于PyTorch的人脸项目也能够看到
这个方案能够部分解决这个问题,但是又会引入新的问题:显存占用不均衡。由于最后结果依然要concat回0号卡上,且loss计算依然在0号卡上,0号卡的显存占用以及计算负载显著高于其他卡。受制于此,依然无法使用较大的batch_size
不仅解决了,还扩展到了更多场景下,支持人脸和re-id训练中常见的margin loss,支持混合精度训练与分布式训练。
几点小小的优势:
batch_size
,训练得更加开心ArcFace
、SphereFace
、CosFace
、AM-softmax
等等首先确认下你是否有必要使用模型并行:
如果以上答案均为肯定,那么你可以考虑使用模型并行。但是由于模型并行需要hack model和optimizer(分布式条件下更为复杂),目前需要自行移植到你的项目中。
原理都是相通的,其他框架如MXNet甚至有对分布式支持更为友好的kvstore
可供使用