Function gen_do_gradient

caffe2/python/control_ops_grad.py:11–166

Generates gradient Do operator, given forward Do op and a list of gradient blobs corresponding to forward op's outputs Returns a gradient op and a list of blobs corresponding to input gradients

(op, g_output)

def gen_do_gradient(op, g_output):
    """
    Generates gradient Do operator, given forward Do op and a list
    of gradient blobs corresponding to forward op's outputs
    Returns a gradient op and a list of blobs corresponding to input gradients
    """
    from caffe2.python.core import BlobReference
    subnet, outer_to_inner_map, inner_to_outer_map, workspace_blob_name = \
        _do_op_sanity_check_and_process(op)

    assert len(g_output) == len(op.output), \
        "Different number of gradient blobs and Do op outputs"

    grad_ops, deduped_g_output = dedupe_g_output(op, g_output)
    g_output = deduped_g_output

    # From the outer net point of view:
    #  Do is an operator that has some number of inputs and outputs;
    #  we have to generate a gradient operator that writes into
    #  corresponding input gradient blobs and has access to inputs, outputs
    #  and gradient output blobs
    # From the inner net point of view:
    #  Do is an operator with a subnet and blob bindings,
    #  we need to forward Do's output blob gradients into inner workspace,
    #  use them to run backward pass generation and forward Do's input blob
    #  gradients back into outer workspace

    op_output = [str(o) for o in op.output]
    op_output = op_output[:-1]  # remove workspace pointer blob
    op_input = [str(i) for i in op.input]
    op_input = op_input[:-1]  # remove workspace pointer blob

    ordered_inner_output_blob_names = [outer_to_inner_map[o] for o in op_output]

    backward_pass_initial_grad_map = {}
    initial_grad_map = {}
    for inner_output_name, outer_grad_output_name in \
            zip(ordered_inner_output_blob_names, g_output):
        # link inner_output_name to corresponding inner_grad_output_name for
        # backward pass generation;
        if outer_grad_output_name:
            inner_grad_output_name = inner_output_name + "/_DO_OPERATOR_INNER_GRAD_"
            backward_pass_initial_grad_map[BlobReference(inner_output_name)] = \
                BlobReference(inner_grad_output_name)
            initial_grad_map[inner_grad_output_name] = str(outer_grad_output_name)
    assert len(initial_grad_map) > 0, "Empty initial gradient map for Do op"

    inner_grad_ops, inner_grad_names_map = _gen_subgradient_pass(
        subnet, backward_pass_initial_grad_map)

    if len(inner_grad_ops) == 0:
        return [], []

    grad_copy_ops = []
    g_input = []
    new_op_outputs = []
    new_blob_bindings = {}
    for outer_input_name in op_input: