Function InjectCrossDeviceCopies

caffe2/python/core.py:2461–2596

Injecting Copy functions between device within a net. Users can provide a net with part of operators using different device_options. This method will automatically create a new net with Copy ops inserted in it. Inputs: blob_to_device: If not None, it is a map of blobs and the

(net, blob_to_device=None, blob_remap=None,
                            placeHolderOps=None)

def InjectCrossDeviceCopies(net, blob_to_device=None, blob_remap=None,
                            placeHolderOps=None):
    '''
    Injecting Copy functions between device within a net. Users can provide
    a net with part of operators using different device_options. This method
    will automatically create a new net with Copy ops inserted in it.

    Inputs:
      blob_to_device: If not None, it is a map of blobs and their device locations.
      blob_remap: If not None, it is a map from a pair (blob, device) to
                  the name of the blob in the given device. Blobs found in this
                  map are assumed to be cached and don't need to be copied.
    Outputs:
      new_net: A new net with CopyCPUToGPU inserted with correct device option

      required_external_to_device:
               A mapping between unresolved external inputs and their
               required device options.
    Assumptions:
      1. every external inputs of this net is already in blob_to_device!
      2. if not, this function will use net device option
      3. InferOpBlobDevices might fail to get the correct inference for ops like
         EnsureCPUOutput that could take in input from multiple places.
    '''
    new_net = net.Clone(net._net.name + '_cross_device', keep_schema=True)
    del new_net._net.op[:]
    if blob_to_device is None:
        blob_to_device = {}
    # remapping of input blobs for each op.
    if blob_remap is None:
        blob_remap = {}
    temp_remap = {}
    net_option = net._net.device_option or caffe2_pb2.DeviceOption()

    # if external_inputs have device remappings generated by previous nets,
    # then add those remappings as external inputs as well.
    all_remaps = defaultdict(list)
    for entry, mapped_blob in blob_remap.items():
        all_remaps[entry.blob].append(mapped_blob)
    mapped_external_inputs = []
    for input in new_net._net.external_input:
        mapped_external_inputs.extend(all_remaps.get(input) or [])
    new_net._net.external_input.extend(mapped_external_inputs)

    for op in net._net.op:
        temp_remap.clear()
        # Get where inputs and outputs should be. If it is a Placeholder
        # (i.e. fake) op, then set op's device as blob's devices.
        input_dev = None
        output_dev = None
        if placeHolderOps is not None and op.type in placeHolderOps:
            input_dev, output_dev = InferOpDeviceAsBlobDevices(op)
        else:
            input_dev, output_dev = InferOpBlobDevices(op)

        for dev, input in zip(input_dev, op.input):
            assert net.BlobIsDefined(input), \
                "input {} should be defined in the net.".format(input)