| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
| from __future__ import unicode_literals |
| |
| import networkx as nx |
| import collections |
| import time |
| import copy |
| from caffe2.python import workspace |
| |
| import logging |
| |
| log = logging.getLogger("memonger") |
| log.setLevel(logging.INFO) |
| LiveRange = collections.namedtuple('LiveRange', ["defined", "used"]) |
| |
| |
| def share_grad_blobs(net, losses, param_grads, namescope): |
| ''' |
| Implements similar optimization as Torch's shareGradInput(): |
| for the gradients that are passed between layers, share blobs between |
| operators when possible. This yields significant memory savings with |
| deep networks. |
| |
| Returns an optimized protobuf (assign to net._net) |
| ''' |
| def is_grad_blob(b): |
| name = str(b) |
| # Note: need to look at _{namescope} pattern as it matches |
| # to handle the auto-split gradients |
| return "_grad" in name and (name.startswith(namescope) or |
| name.startswith("_" + namescope)) and name not in param_grads |
| |
| def is_grad_op(op): |
| # TODO: something smarter |
| for inp in op.input: |
| if is_grad_blob(inp): |
| return True |
| for out in op.output: |
| if is_grad_blob(out): |
| return True |
| return False |
| |
| start_time = time.time() |
| log.warn("NOTE: Executing *experimental* memonger to " + |
| "optimize gradient memory") |
| |
| # Collect ops that have something to do with |
| # gradients |
| if not namescope.endswith("/"): |
| namescope += "/" |
| |
| netproto = copy.deepcopy(net.Proto()) |
| grad_ops = [op for op in netproto.op if is_grad_op(op)] |
| |
| # Create mapping from blobs to ops |
| blobs_to_ops = collections.defaultdict(lambda: []) |
| blob_input_count = collections.defaultdict(lambda: 0) |
| op_inputs = collections.defaultdict(lambda: 0) |
| op_visit_count = collections.defaultdict(lambda: 0) |
| for i, op in enumerate(grad_ops): |
| for inp in op.input: |
| if is_grad_blob(inp) or inp in losses: |
| # Ignore in-place transformation ops (self cycles) |
| if inp not in op.output: |
| blobs_to_ops[inp].append(i) |
| op_inputs[i] += 1 |
| |
| # Traverse operators starting from the loss blobs. |
| # Keep tabs on when blobs are seen first and last, and also |
| # when operators have their input satisfied. Share blobs only |
| # under same branch, avoiding problems with parallel workers. |
| output_blobs = set() |
| mapping = {} |
| |
| def descend(op_idx, free_blobs): |
| cur_op = grad_ops[op_idx] |
| new_free_blobs = set() |
| for inp in cur_op.input: |
| if is_grad_blob(inp): |
| blob_input_count[inp] += 1 |
| if blob_input_count[inp] == len(blobs_to_ops[inp]): |
| actual_blob = inp if inp not in mapping else mapping[inp] |
| new_free_blobs.add(actual_blob) |
| |
| for outp in cur_op.output: |
| if is_grad_blob(outp): |
| if outp not in output_blobs: |
| # First seen this blob as output, can assign to a free blob |
| for freeb in free_blobs: |
| mapping[outp] = freeb |
| free_blobs.remove(freeb) |
| break |
| |
| output_blobs.add(outp) |
| |
| free_blobs.update(new_free_blobs) |
| |
| first_branch = True |
| for outp in cur_op.output: |
| for inp_op_idx in blobs_to_ops[outp]: |
| op_visit_count[inp_op_idx] += 1 |
| |
| # Descend only if we have satisfied all inputs |
| if op_visit_count[inp_op_idx] == op_inputs[inp_op_idx]: |
| free_blobs_fwd = free_blobs if first_branch else set() |
| first_branch = False |
| descend(inp_op_idx, free_blobs_fwd) |
| |
| # Start DFS from the losses |
| for loss in losses: |
| for op_idx in blobs_to_ops[loss]: |
| descend(op_idx, set()) |
| |
| # Rename the shared blobs |
| shared_blobs = set(mapping.values()) |
| renamed = {} |
| for j, b in enumerate(shared_blobs): |
| renamed[b] = namescope + "__m{}_".format(j) |
| |
| # Final mapping |
| for k, v in mapping.items(): |
| mapping[k] = renamed[v] |
| |
| # Add the originators |
| mapping.update(renamed) |
| log.info("Remapping {} blobs, using {} shared".format( |
| len(mapping), len(renamed), |
| )) |
| apply_assignments(netproto, mapping) |
| log.info("Gradient memory optimization took {} secs".format( |
| time.time() - start_time), |
| ) |
| return netproto |
| |
| |
| def topological_sort_traversal(g): |
| return nx.topological_sort(g) |
| |
| |
| def compute_ranges(linearized_ops): |
| blobs = collections.defaultdict(lambda: LiveRange(defined=None, used=None)) |
| for i, op in enumerate(linearized_ops): |
| for blob in op.input: |
| used = blobs[blob].used |
| if used is None: |
| used = i |
| else: |
| used = max(used, i) |
| blobs[blob] = blobs[blob]._replace(used=used) |
| for blob in op.output: |
| defined = blobs[blob].defined |
| if defined is None: |
| defined = i |
| else: |
| defined = min(defined, i) |
| blobs[blob] = blobs[blob]._replace(defined=defined) |
| |
| return blobs |
| |
| |
| def is_compatible(candidate_range, assignment, static_blobs): |
| (name, range_) = assignment[-1] |
| if name in static_blobs: |
| return False |
| if candidate_range.defined is None or range_.defined is None \ |
| or range_.used is None: |
| return False |
| return candidate_range.defined > range_.used |
| |
| |
| def compute_blob_assignments(assignments): |
| blob_assignments = {} |
| for assignment in assignments: |
| if len(assignment) == 1: |
| continue |
| last_blob, _ = assignment[-1] |
| for (blob, _) in assignment: |
| blob_assignments[blob] = last_blob |
| return blob_assignments |
| |
| |
| def compute_assignments(ranges, static_blobs): |
| # Sort the ranges based on when they are last used. |
| # If LiveRange.used is None, then the blob is never used and could |
| # be consumed externally. Sort these to the end of the list as opposed |
| # to the beginning so that they can be shared as well. |
| ranges = sorted( |
| list(ranges.items()), |
| key=lambda p: (p[1].used is None, p[1].used), |
| ) |
| assignments = [] |
| for (name, range_) in ranges: |
| assigned = False |
| for assignment in assignments: |
| if is_compatible(range_, assignment, static_blobs): |
| assignment.append((name, range_)) |
| assigned = True |
| break |
| if assigned: |
| continue |
| assignments.append([(name, range_)]) |
| return assignments |
| |
| |
| def compute_interference_graph(ops): |
| g = nx.DiGraph() |
| for i, op in enumerate(ops): |
| g.add_node(i, op=op) |
| for i, parent_op in enumerate(ops): |
| for j, child_op in enumerate(ops): |
| if i == j: |
| continue |
| if any(output in child_op.input for output in parent_op.output): |
| deps = set(child_op.input).intersection(parent_op.output) |
| g.add_edge(i, j, deps=deps) |
| assert nx.is_directed_acyclic_graph(g), child_op |
| return g |
| |
| |
| Optimization = collections.namedtuple( |
| 'Optimization', ['net', 'assignments', 'blob_assignments']) |
| |
| |
| def apply_assignments(net, blob_assignments): |
| def canonical_name(blob): |
| if blob not in blob_assignments: |
| return blob |
| return blob_assignments[blob] |
| |
| for op in net.op: |
| for i, input_ in enumerate(op.input): |
| op.input[i] = canonical_name(input_) |
| for i, output in enumerate(op.output): |
| op.output[i] = canonical_name(output) |
| |
| |
| def optimize_interference(net, static_blobs, |
| ordering_function=topological_sort_traversal): |
| """ |
| 1) Use a BFS traversal of the execution graph to generate an |
| ordering of the node executions. |
| 2) Generate use-def ranges for each `blob` in the BFS traversal |
| order. |
| 3) Assign blobs to `canonical blobs` |
| 4) Rename blobs to canonical blobs |
| """ |
| net = copy.deepcopy(net) |
| g = compute_interference_graph(net.op) |
| ordering = ordering_function(g) |
| linearized_ops = [net.op[i] for i in ordering] |
| |
| # Reorder ops in net based on the computed linearlized order. |
| # If the graph has multiple topological orderings and if the NetDef's |
| # ordering differs from the order used to compute ranges, then the |
| # runtime might end up overwriting blobs before they are used. |
| del net.op[:] |
| net.op.extend(linearized_ops) |
| |
| ranges = compute_ranges(linearized_ops) |
| assignments = compute_assignments(ranges, static_blobs) |
| blob_assignments = compute_blob_assignments(assignments) |
| apply_assignments(net, blob_assignments) |
| return Optimization( |
| net=net, |
| blob_assignments=blob_assignments, |
| assignments=assignments) |
| |
| Statistics = collections.namedtuple( |
| 'Statistics', ['baseline_nbytes', 'optimized_nbytes']) |
| |
| |
| def compute_statistics(assignments): |
| def blob_nbytes(blob): |
| return workspace.FetchBlob(blob).nbytes |
| blob_bytes = { |
| blob: blob_nbytes(blob) for assignment in assignments |
| for (blob, _) in assignment} |
| baseline_nbytes = sum(v for _, v in blob_bytes.iteritems()) |
| optimized_nbytes = sum( |
| max(blob_bytes[blob] for (blob, _) in assignment) |
| for assignment in assignments) |
| return Statistics( |
| baseline_nbytes=baseline_nbytes, |
| optimized_nbytes=optimized_nbytes) |
