Pytorch多GPU和Sync BatchNorm代码

由于复现spade的过程中遇到了一些GPU的问题，所以决定好好理解一下DPL

代码问题

终端显示暂行不动，并没有报错。GPU没有加载进程，同时CPU也没有动。
考虑可能是DPL的问题。由于代码中使用了Sync BatchNorm，考虑到可能是DPL的问题。

nn.DataParallel

在forward阶段，当前GPU上的module会被复制到其他GPU上，输入数据则会被切分，分别传到不同的GPU上进行计算；在backward阶段，每个GPU上的梯度会被求和并传回当前GPU上，并更新参数。也就是复制module -> forward -> 计算loss -> backward -> 汇总gradients -> 更新参数 -> 复制module -> …的不断重复执行，示意图如下：
因为数据会被均分到不同的GPU上，所以要求batch_size大于GPU的数量。下面对DataParallel的forward函数做一个简单的解释：

class DataParallel(Module):
    def __init__(self, module, device_ids=None, output_device=None, dim=0):
        super(DataParallel, self).__init__()
        if not torch.cuda.is_available():
            self.module = module
            self.device_ids = []
            return
        if device_ids is None:
            device_ids = list(range(torch.cuda.device_count()))
        if output_device is None:
            output_device = device_ids[0]
        self.dim = dim
        self.module = module   # 待并行计算的模型
        self.device_ids = list(map(lambda x: _get_device_index(x, True),device_ids))
        self.output_device = _get_device_index(output_device, True)
        self.src_device_obj = torch.device("cuda:{}".format(self.device_ids[0]))
        _check_balance(self.device_ids)
        if len(self.device_ids) == 1:
            self.module.cuda(device_ids[0])

    def forward(self, *inputs, **kwargs):
        if not self.device_ids:
            return self.module(*inputs, **kwargs)
        for t in chain(self.module.parameters(), self.module.buffers()):
            if t.device != self.src_device_obj:
                raise RuntimeError("module must have its parameters and buffers "
                                   "on device {} (device_ids[0]) but found one of "
                                   "them on device: {}".format(self.src_device_obj, t.device))

        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)  
        if len(self.device_ids) == 1:
            return self.module(*inputs[0], **kwargs[0])
        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
        outputs = self.parallel_apply(replicas, inputs, kwargs)
        return self.gather(outputs, self.output_device)

    def scatter(self, inputs, kwargs, device_ids):
        '''scatter_kwargs内部调用名为scatter的函数，
        作用是将输入数据及参数均分到每张卡上，以及复制其他类型对象的引用'''
        return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)

    def replicate(self, module, device_ids):
        '''replicate对输入模型的parameters、buffers、modules都一一进行copy，并返回copy的list，
        因为modules最终是以类似链表的形式存储的，所以list中只包含第一个module'''
        return replicate(module, device_ids)
        
    def parallel_apply(self, replicas, inputs, kwargs):
        '''内部调用python的Thread将分割好的input分配到不同的GPU上计算，并返回result dict
        这里会调用SyncBN的前向反馈方法'''
        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])

    def gather(self, outputs, output_device):
        '''将每张卡的计算结果统一汇聚到主卡，从不同GPU上取回结果'''
        return gather(outputs, output_device, dim=self.dim)

原生的该方法只是将模型在每张卡上复制一份，并且没有建立起联系，而我们的 SyncBN 是需要进行同步的，因此需要重载该方法，让各张卡上的SyncBN 通过某种数据结构和同步机制建立起联系。

重载nn.DataParallel.replicate方法

可以设计一个继承nn.DataParallel的子类DataParallelWithCallBack，重载了replicate方法，子类的该方法先是调用父类的replicate方法，然后调用一个自定义的回调函数（这也是之所以命名为DataParallelWithCallBack的原因），该回调函数用于将各卡对应的 SyncBN 层关联起来，使得它们可以通过某种数据结构进行通信。

class DataParallelWithCallback(DataParallel):
    def replicate(self, module, device_ids):
        modules = super(DataParallelWithCallback, self).replicate(module, device_ids)
        execute_replication_callbacks(modules)#回调函数
        return modules

子类重载的replicate方法

def execute_replication_callbacks(modules):
    master_copy = modules[0]#主卡上的模型
    nr_modules = len(list(master_copy.modules()))
    ctxs = [CallbackContext() for _ in range(nr_modules)]
    #上下文数据结构，用于将各卡上对应的SynBN层关联起来
    for i, module in enumerate(modules):
        for j, m in enumerate(module.modules()):
            if hasattr(m, '__data_parallel_replicate__'):#判断SynBN层的条件
                m.__data_parallel_replicate__(ctxs[j], i)#其余各卡向主卡进行注册关联

自定义的回调函数，将各卡对应的Syn-BN层进行关联，其中CallbackContext是一个自定义类，其中没有定义实质性的东西，作为一个上下文数据结构，实例化这个类的对象主要用于将各个卡上对应的Syn-BN层进行关联；__data_parallel_replicate__是在Syn-BN中定义的方法，在该方法中其余子卡上的Syn-BN层会向主卡进行注册，使得主卡能够通过某种数据结构和各卡进行通信。

SynBN的同步注册机制

由上可知，我们需要在 SynBN 中实现一个用于同步的注册方法，SynBN 中还需要设置一个用于管理同步的对象（下图中的 _sync_master），这个对象有一个注册方法，可将子卡注册到其主卡。

在 SyncBN 的方法中，若是主卡，则将上下文管理器的 sync_master 属性设置为这个管理同步的对象（_sync_master）；否则，则调用上下文对象的同步管理对象的注册方法，将该卡向其主卡进行注册。

    def __data_parallel_replicate__(self, ctx, copy_id):
        self._is_parallel = True
        self._parallel_id = copy_id
        # parallel_id == 0 means master device.
        if self._parallel_id == 0:
            ctx.sync_master = self._sync_master
        else:
            self._slave_pipe = ctx.sync_master.register_slave(copy_id)#若非主卡则进行注册

SynBN的同步注册机制

    def register_slave(self, identifier):
        if self._activated:
            assert self._queue.empty(), 'Queue is not clean before next initialization.'
            self._queue.queue.clear()#用于主卡和各子卡之间进行数据共享
            self._activated = False
            self._registry.clear()
        future = FutureResult()#存储各子卡的计算结果
        self._registry[identifier] = _MasterRegistry(future)#注册子卡
        return SlavePipe(identifier, self._queue, future)#返回一个子卡用于同步的对象

主卡进行同步管理的类中注册子卡的方法

class SyncMaster(object):#主卡通过实例化该对象来进行同步管理
    def __init__(self, master_callback):
        self._master_callback = master_callback
        self._queue = queue.Queue()#主卡与子卡共享队列传递与获取计算结果
        self._registry = collections.OrderedDict()#通过它识别各子卡
        self._activated = False

主卡进行同步管理的类

_MasterRegistry = collections.namedtuple('MasterRegistry', ['result'])
_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result'])
class SlavePipe(_SlavePipeBase):
    def run_slave(self, msg):
        self.queue.put((self.identifier, msg))
        ret = self.result.get()
        self.queue.put(True)
        return ret

子卡进行同步操作的类

SynBN的前向反馈

首先，每张卡上的 SyncBN 各自计算出 mini-batch 的均值和与平方和，然后主卡上的 SyncBN 收集来自各个子卡的计算结果，从而计算出全局的均值和方差，接着发放回各个子卡，最后各子卡的 SyncBN 收到来自主卡返回的计算结果各自进行归一化（和缩放平移）操作。当然，主卡上的 SyncBN 计算出全局统计量后就可以进行它的归一化（和缩放平移）操作了。