Function sumN

caffe2/python/data_parallel_model.py:1054–1077

Create a Sum op for 2 or more blobs on different devices. Saves the result on the first device. Args: dev_indices -- a list of device indices, which can be translated into CUDA identifiers with model._devices

(*dev_indices)

        p2p_access_pattern = None

    def sumN(*dev_indices):
        """Create a Sum op for 2 or more blobs on different devices.
        Saves the result on the first device.

        Args:
        dev_indices -- a list of device indices, which can be translated into
                       CUDA identifiers with model._devices
        """
        devices = [model._devices[idx] for idx in dev_indices]
        blobs = [blobs_group[idx] for idx in dev_indices]
        device_opt = core.DeviceOption(model._device_type, devices[0])
        with core.DeviceScope(device_opt):
            for i, peer in enumerate(devices):
                if i == 0:
                    continue  # Skip the first device
                if p2p_access_pattern is not None and p2p_access_pattern.size and not p2p_access_pattern[
                    devices[0], peer
                ]:
                    # Copy from peer to d0
                    blobs[i] = model.Copy(
                        blobs[i],
                        'gpu_{}/{}_gpu{}_copy'.format(devices[0], param, peer)
                    )
            net.Sum(blobs, [blobs[0]], name='dpm')

    if len(devices) == 16:
        # Special tree reduction for 16 gpus, TODO generalize like in muji.py