mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-26 13:18:31 +08:00
PERF: Release unneeded pipeline terms.
Refcount pipeline terms during execution and release terms once they're no longer needed. This dramatically reduces memory usage on large pipelines.
This commit is contained in:
@@ -344,6 +344,8 @@ class SimplePipelineEngine(object):
|
||||
loader_group_key = juxt(get_loader, getitem(graph.extra_rows))
|
||||
loader_groups = groupby(loader_group_key, graph.loadable_terms)
|
||||
|
||||
refcounts = graph.initial_refcounts()
|
||||
|
||||
for term in graph.ordered():
|
||||
# `term` may have been supplied in `initial_workspace`, and in the
|
||||
# future we may pre-compute loadable terms coming from the same
|
||||
@@ -380,6 +382,14 @@ class SimplePipelineEngine(object):
|
||||
else:
|
||||
assert workspace[term].shape == (mask.shape[0], 1)
|
||||
|
||||
# Decref any term we depended on.
|
||||
for (parent, _) in graph.in_edges(term):
|
||||
refcounts[parent] -= 1
|
||||
# No one else depends on this term. Remove it from the
|
||||
# workspace to conserve memory.
|
||||
if refcounts[parent] == 0:
|
||||
del workspace[parent]
|
||||
|
||||
out = {}
|
||||
graph_extra_rows = graph.extra_rows
|
||||
for name, term in iteritems(graph.outputs):
|
||||
|
||||
@@ -119,6 +119,19 @@ class TermGraph(DiGraph):
|
||||
def _repr_png_(self):
|
||||
return self.png.data
|
||||
|
||||
def initial_refcounts(self):
|
||||
"""
|
||||
Calculate initial refcounts for execution of this graph.
|
||||
|
||||
Each node starts with a refcount equal to its outdegree, and output
|
||||
nodes get one extra reference to ensure that they're still in the graph
|
||||
at the end of execution.
|
||||
"""
|
||||
refcounts = self.out_degree()
|
||||
for t in self.outputs.values():
|
||||
refcounts[t] += 1
|
||||
return refcounts
|
||||
|
||||
|
||||
class ExecutionPlan(TermGraph):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user