Function _ComputeBlobsToSync

caffe2/python/data_parallel_model.py:1691–1731

We sync all blobs that are generated by param init net and are 'data parallel', i.e assigned to a device

(model)

def _ComputeBlobsToSync(model):
    '''
    We sync all blobs that are generated by param init net and
    are 'data parallel', i.e assigned to a device
    '''
    sync_names = set()

    # We don't sync params if the model is shared
    if model._shared_model:
        blobs_to_sync = [str(p) for p in model.GetComputedParams('')]
        sync_names = [stripBlobName(p) for p in blobs_to_sync]
    else:
        blobs_to_sync = []

        for op in model.param_init_net.Proto().op:
            dp_outputs = [
                o for o in op.output
                if o.startswith("{}_".format(model._device_prefix))
            ]
            sync_names.update([stripBlobName(o) for o in dp_outputs])
            blobs_to_sync.extend(dp_outputs)

        # Sanity check
        diff = set(model._param_names) - sync_names
        assert diff == set(), \
           "Some params not instantiated in param init net: {}".format(diff)

    # Remove duplicates and sort
    prefixlen = len(model._device_prefix) + 1

    def extract_sort_key(b):
        # Sort first based on device id, and then by whole string
        deviceid = int(b[prefixlen:b.index(scope._NAMESCOPE_SEPARATOR)])
        return (deviceid, b)

    blobs_to_sync = sorted(
        list(set(blobs_to_sync)),
        key=extract_sort_key)

    blobs_to_sync = [core.BlobReference(b) for b in blobs_to_sync]
    return (blobs_to_sync, sync_names)


def _OptimizeGradientMemorySimple(model, losses_by_gpu, devices):