Make metapath walk faster by grouping adjacency lists by node type #1399
Conversation
|
Code Climate has analyzed commit 1e4d782 and detected 3 issues on this pull request. Here's the issue category breakdown:
View more on Code Climate. |
|
Unfortunately this seems to duplicate a bunch of the node-data/edge-data interaction work from #1220, unlucky! 🙁 A 29× speed-up is pretty nice!
As a potential tweak to the implementation of this, maybe we could take inspiration from #1291 and #1316, expanding
This has the downside of requiring multiple batches of memory, but it may be the right trade-off. In particular, I'm assuming that the workflow would typically be an interactive session with small-ish graphs for experimenting with different algorithms (thus requiring multiple indexing schemes), where the user has more control and monitoring of memory usage. A real ML pipeline on larger data would then be only using one or two algorithms, and thus only creating and storing the indices that are really needed. What do you think? |
That's an interesting idea. So we would build the index automatically when
As an alternative, maybe it could be specific generators/walkers that force a graph to build the index, and default to the original behaviour of filtering the adjacency list after obtaining them if the index doesn't exist class StellarGraph:
def build_adjacency_index(index_by)
# do nothing if index exists otherwise build
class UniformRandomMetapathWalk:
def __init__(self, graph, ...):
self._graph = graph
def run(self, ...):
self._graph.build_adjacency_index("other_node_type")or something like that? |
|
nice speedup! Unfortunately I think part of this speedup is due to not having to pay the cost of converting the node ids to ilocs in neighbour_types = self.graph.node_type(neighbours, use_ilocs=True)
neighbours = [
n_node
for n_node, n_type in zip(neighbours, neighbour_types)
if n_type == metapath[d]
]gets another 2x speedup, so 15x in total but still shy of your 29x! But maybe this could be a quick way of speeding up That being said, I think this is still a valuable way of storing the edges, and we should definitely implement this as an edge caching option! |
|
On the topic of edge caching, for unweighted graphs I think we can get a ~20x speedup (on top of this 29x speedup) if we store the In the case of |
|
@kjun9 can you run the allocation benchmarks for this PR? |
Yeah.
I guess that's equivalent in my mind to
If you've done that work and measured it, can we have a PR for it?
I think that
Typo: "numpy lookups"?
I think there's enough above that we can potentially pause effort on this PR for a bit (at least until the other PRs land) and not spent too much time doing comparisons with it until we've done that baseline work, since those comparisons will be invalidated anyway? |
I was thinking that building/not-building an index "explicitly" would potentially be useful e.g. an algorithm just requires a single call to I've adjusted the code to do the building lazily, taking inspiration from #1291 and also thinking a little bit about the extensions I'll include the new benchmarks here anyway but we can keep in mind the comparisons will be invalidated by the other PRs (still a bit slower due to additional if statements, but not quite as bad, and the allocation was probably much worse before): |
There was a problem hiding this comment.
This is nice. I think we can even extend this idea to handle edge types too, and so get optimised and memory use queries for HinSAGE (since that currently constructs its own type_adjacency_list thing?), and for something like #433.
stellargraph/core/element_data.py
Outdated
| def _numpyise_inner(d, dtype): | ||
| return {k: _numpyise(v, dtype) for k, v in d.items()} | ||
|
|
||
| self._edges_index_by_other_node_type = _AdjIndex( |
There was a problem hiding this comment.
This ends up being a dictionary like {node_id: {node_type: array}}, right?
stellargraph/core/element_data.py
Outdated
| in_dict.setdefault(tgt, {}).setdefault(src_type, []).append(i) | ||
| out_dict.setdefault(src, {}).setdefault(tgt_type, []).append(i) | ||
|
|
||
| undirected.setdefault(tgt, {}).setdefault(src_type, []).append(i) |
There was a problem hiding this comment.
Also, I wonder if it actually makes sense to compute each of these separately (following #1316)? E.g. an algorithm that's using undirected probably won't touch in_dict/out_dict, and vice versa. I thing this might be especially important for this, because there'll be a lot more memory overhead from having many smaller np.ndarray instances (i.e. there'll be more so more of the fixed overheads of each one).
There was a problem hiding this comment.
Hmm yeah makes sense. I might try to extend the _AdjIndex type to make it lazy for individual dicts
|
I think this now conflicts a lot with #1435 and the new speedups I'm getting are not as huuuuge compared to what's already landed in |
|
A 3× speedup still seems worth it, but yeah, maybe we can leave it till after the 1.0 crunch. |
|
Hmm maybe the benchmark test doesn't quite provide the full picture.. notebook seems to be much slower than 3x (been running random walks for ~2 hours) |
|
Hm, 2 hours with this PR or with |
with develop. With this PR, the whole notebook runs within ~10 mins or so |
stellargraph/data/explorer.py
Outdated
|
|
||
| walks.append( | ||
| list(self.graph.node_ilocs_to_ids(walk)) | ||
| ) # store the walk | ||
|
|
||
| return walks | ||
|
|
||
| def _augment_and_group_metapaths(self, metapaths, walk_length): |
There was a problem hiding this comment.
What does "augment" mean here? Maybe:
| def _augment_and_group_metapaths(self, metapaths, walk_length): | |
| def _extend_and_group_metapaths(self, metapaths, walk_length): |
or even
| def _augment_and_group_metapaths(self, metapaths, walk_length): | |
| def _repeat_and_group_metapaths(self, metapaths, walk_length): |
There was a problem hiding this comment.
extend sounds good to me. augment came from me taking the phrasing from an existing comment that I removed # augment metapath to be length long.
tests/core/test_graph.py
Outdated
| @@ -762,7 +762,7 @@ def f(): | |||
| @pytest.mark.parametrize("num_nodes,num_edges", [(0, 0), (100, 200), (1000, 5000)]) | |||
| # features or not, to capture their cost | |||
| @pytest.mark.parametrize("feature_size", [None, 100]) | |||
| @pytest.mark.parametrize("force_adj_lists", [None, "directed", "undirected", "both"]) | |||
| @pytest.mark.parametrize("force_adj_lists", [None, "homogeneous", "by_other_node_type"]) | |||
There was a problem hiding this comment.
This is dramatically changing what this benchmark is measuring; previously it was capturing the cost of using a graph in undirected mode or directed mode (or neither (adjacency matrix only), or both), now it's only capturing the cost of using a graph in either:
- adjacency matrix only
- both undirected and directed mode for homogeneous
- both undirected and directed mode for other_node_type
This benchmark is kinda completely exploding in terms of options, but maybe it instead makes sense to do this as a new parameter like:
@pytest.mark.parametrize("adj_list_directness", ["homogeneous", "by_other_node_type"])Potentially with a:
if force_adj_lists is None and adj_list_directness != "homogeneous":
pytest.skip("only measure force_adj_lists=None once")However, maybe it's not entirely useful to know the specific cost of the undirected vs. directed adjacencies.
tests/data/test_metapath_walker.py
Outdated
| @@ -317,11 +317,11 @@ def test_init_parameters(self): | |||
| assert all(np.array_equal(w1, w2) for w1, w2 in zip(run_1, run_2)) | |||
|
|
|||
| def test_benchmark_uniformrandommetapathwalk(self, benchmark): | |||
| g = example_graph_random(n_nodes=50, n_edges=500, node_types=2) | |||
| g = example_graph_random(n_nodes=50, n_edges=2000, node_types=4) | |||
There was a problem hiding this comment.
Out of curiosity, why did this need to become larger? The metapaths got truncated often in practice, despite the theory?
There was a problem hiding this comment.
I tried increasing it since I wasn't so convinced that the speedups reported from the benchmark were as good as what we end up seeing in practice (#1399 (comment) vs #1399 (comment) the larger benchmark reports much higher speedup)
There was a problem hiding this comment.
Is that the number of edges increasing or the number of head nodes increasing?
tests/data/test_metapath_walker.py
Outdated
| mrw = UniformRandomMetaPathWalk(g) | ||
|
|
||
| # this should be made larger to be more realistic, when it is fast enough | ||
| nodes = np.arange(0, 5) | ||
| nodes = np.arange(0, 10) |
There was a problem hiding this comment.
This is probably now fast enough to switch to being the same as the other ones, e.g.
| nodes = np.arange(0, 10) | |
| nodes = np.arange(0, 50) |
stellargraph/core/element_data.py
Outdated
| self.undirected = self.init_index(ins=True, outs=True) | ||
| self.ins = self.init_index(ins=True, outs=False) | ||
| self.outs = self.init_index(ins=False, outs=True) |
There was a problem hiding this comment.
Rather than duplicate the init_index calling and assignment logic, maybe this could be:
| self.undirected = self.init_index(ins=True, outs=True) | |
| self.ins = self.init_index(ins=True, outs=False) | |
| self.outs = self.init_index(ins=False, outs=True) | |
| self.lookup(ins=True, outs=True) | |
| self.lookup(ins=True, outs=False) | |
| self.lookup(ins=False, outs=True) |
stellargraph/core/element_data.py
Outdated
| @@ -303,74 +303,117 @@ def _numpyise(d, dtype): | |||
| return {k: np.array(v, dtype=dtype) for k, v in d.items()} | |||
|
|
|||
|
|
|||
| class AdjList: | |||
| def __init__(self, init_index): | |||
| self.init_index = init_index | |||
There was a problem hiding this comment.
Could you talk a bit about the design choice of passing in a function here? I'm not saying it's wrong and I'm not sure there's a best answer at all, but I wonder if there's other versions that are a little clearer?
In particular, it feels a bit funny to have EdgeData store an AdjList object that holds a method from that parent EdgeData. An alternative might be something like
class AdjList:
def __init__(self, by_other_node_type, ...):
if by_other_node_type:
self._init_func = self._init_by_other_node_type
else:
self._init_func = self._init_plain
self.undirected = self.ins = self.outs = None
def _init_plain(self, ...):
...
def _init_by_other_node_type(self, ...):
...I guess we'll need to pass more arguments around in this case.
There was a problem hiding this comment.
Hmm yeah I wanted AdjList to handle the lookup + lazy initialising in one place, but thought it would be more convenient to define more adjacency-initialiser functions in EdgeData rather than in AdjList since I'm guessing it'll need access to all the type information in EdgeData and NodeData to build the adjacency lists with, so that's how I ended up with this approach.
Not sure if you had something else in mind with your suggestion, but I guess we could pass in the whole edge data and node data into the adjacency list to do whatever it needs to do... but maybe that also looks a bit funny?
class AdjList:
def __init__(self, by_other_node_type, edge_data, node_data)
def _init_plain(self, ...):
def _init_by_other_node_type(self, ...):
class EdgeData:
def __init__(self, ...):
...
self._adj_list = AdjList(by_other_node_type=False, self, node_data)
self._adj_list_by_other_node_type = AdjList(by_other_node_type=True, self, node_data)There was a problem hiding this comment.
Yeah, that's a bit funny too.
I thought about this overnight and I think my "concern" here is that the recursiveness of the data structure initialisation introduces risks of attempting to use a value while it's being initialised (e.g. if edge list initialisation happened to look at other edge lists, for some reason).
I guess the fully-separated version would be passing in sources, targets and node_data?
However, this may not be worth it, especially as edge list initialisation is:
- limited enough that there doesn't seem to be a reason why it would look at itself
- done as a final commit, rather than incremental construction (specifically, the
self.undirectedproperties go fromNoneto "fully ready", rather than being incrementally built up), and so we'll see infinite recursion here rather than invalid results
The main risk might be around delaying garbage collection (since this causes a cycle between objects, and so cPython will have to wait for the actual garbage collector to run, rather than objects being deallocated as soon as they go out of scope), and future refactoring that change either of the points above.
stellargraph/core/element_data.py
Outdated
| super().__init__(shared, type_starts) | ||
|
|
||
| # cache these columns to avoid having to do more method and dict look-ups | ||
| self.sources = self._column(SOURCE) | ||
| self.targets = self._column(TARGET) | ||
| self.weights = self._column(WEIGHT) | ||
|
|
||
| self.source_types = nodes.type_of_iloc(self.sources) |
There was a problem hiding this comment.
I think this is going to noticeably increase memory usage, both because:
- it's using type names not type ilocs
- it's two extra O(num edges) arrays that aren't shared with anything (i.e. they're new) but are stored permanently
The first can be fixed with:
| self.source_types = nodes.type_of_iloc(self.sources) | |
| self.source_types = nodes.type_ilocs[self.sources] |
but the second is more problematic. Maybe these could be computed lazily too? (And/or, not even stored, just used transiently for initialising the adjacency lists?)
I think the benchmarks in the PR description aren't up to date with the current state? Maybe some new benchmarks could be reported, including the test_allocation_benchmark ones (to try to capture this addition)?
There was a problem hiding this comment.
Good point, yeah I don't think we need to store these at all actually. I'll run some new benchmarks afterwards
stellargraph/core/element_data.py
Outdated
| @@ -303,74 +303,117 @@ def _numpyise(d, dtype): | |||
| return {k: np.array(v, dtype=dtype) for k, v in d.items()} | |||
|
|
|||
|
|
|||
| class AdjList: | |||
There was a problem hiding this comment.
It's minor, but I think this isn't quite an adjacency list, but actually something like:
| class AdjList: | |
| class LazyAdjLists: |
This reverts commit 1ce0976.
|
Update: I decided to kind of start-over by getting rid of the Adjacency list construction times (should be roughly the same): Adjacency list allocations: Metapath walk (I ran the test in this branch against the code in develop, since I tweaked the test slightly to look at how times change as the graph gets larger/more heterogeneous). Looks like the speedup roughly ranges between 2x - 3x? |
|
I've switched to contiguous arrays like #1296 - took me a little while to understand how to do it.. 😅 I did try to take out a few small functions to share with the existing code for constructing the adjacency lists, but New benchmarks after switching to contiguous arrays: Metapath2vec walkAdjacency list construction timesAdjacency list allocations |
There was a problem hiding this comment.
This looks awesome. It's definitely much harder for me to think about this implementation, but it looks like it's not too much slower than the non-node type one, and (based on comparing the benchmarks in the PR description with your latest ones) significantly faster than the old form. 👍
| neigh_counts = sum( | ||
| np.bincount(ilocs, minlength=number_of_nodes) for ilocs in node_ilocs_list | ||
| ) |
There was a problem hiding this comment.
I think this will allocate len(node_iloc_list) extra arrays, because it's not mutating in place like the old form did (in particular it'll do something like (0 + bincount) + bincount, and each of those + will create a whole new array). We could resolve this with something more complicated like:
| neigh_counts = sum( | |
| np.bincount(ilocs, minlength=number_of_nodes) for ilocs in node_ilocs_list | |
| ) | |
| neigh_counts = np.bincount(node_ilocs_list[0], minlength=number_of_nodes) | |
| for ilocs in node_ilocs_list[1:]: | |
| neigh_counts += np.bincount(ilocs, minlength=number_of_nodes) |
But, based on the benchmarks not really changing, I don't think this is worth it, especially because this would be entirely better using numpy/numpy#7998.
| source_types = self.node_data.type_ilocs[self.sources] | ||
| target_types = self.node_data.type_ilocs[self.targets] | ||
|
|
||
| def _to_directed_adj_list(arr, other_node_types): |
There was a problem hiding this comment.
The peak memory use may be able to be reduced by creating the ..._types arrays more transiently, so that they don't both exist at once. In particular, I think they can probably moved inside the helper method:
| source_types = self.node_data.type_ilocs[self.sources] | |
| target_types = self.node_data.type_ilocs[self.targets] | |
| def _to_directed_adj_list(arr, other_node_types): | |
| def _to_directed_adj_list(node_ilocs, other_node_ilocs): | |
| other_node_types = self.node_data.type_ilocs[other_node_ilocs] |
And then call it like _to_directed_adj_list(self.targets, self.sources) (or vice versa) below.
| source_types = self.node_data.type_ilocs[self.sources] | ||
| target_types = self.node_data.type_ilocs[self.targets] | ||
|
|
||
| def _to_directed_adj_list(arr, other_node_types): |
There was a problem hiding this comment.
If you felt like it, this nested method could become a top level method, like:
def _to_directed_adj_list(self, node_ilocs, other_node_ilocs):
...
def _create_directed_adj_lists_by_other_node_type(self):
return (
self._to_directed_adj_list(self.targets, self.sources),
self._to_directed_adj_list(self.sources, self.targets),
)| for other_node_type in np.unique(other_node_types): | ||
|
|
||
| # filter node ilocs based on other node type | ||
| arr_filtered = arr[other_node_types == other_node_type] |
There was a problem hiding this comment.
Because of the various whole-array filters here, this loop is O(edges * node types + nodes * node types); it's probably fine, assuming there's not many node types, but we could potentially consider other approaches, e.g.
- doing a Pandas groupby, which might be
O(edges + nodes * node types), although it'll like have more overhead, e.g.for other_node_type, flat_filtered in pd.Series(flat).groupby(other_node_types): splits = _get_splits([arr[flat_filtered]], self.number_of_nodes, self._id_index.dtype) index[other_node_type] = FlatAdjacencyList(flat_filtered, splits)
- sorting
other_node_types(and other bits too) so that index returned bynp.unique(other_node_types, return_index=True)gives the first index of each node type, and thus we can do O(1) slicing instead of O(edges) comparisons
Both of these are likely to have more overhead and potentially be a bit fiddly (I don't know if the first one is the right code), so I don't know if they'll make sense. I don't know if we have a good "real world" dataset with many node types to test this, but this may also indicate that it's not actually a concern (knowledge graphs can end up with thousands of edge types... but only one node type).
Thoughts?
There was a problem hiding this comment.
The pandas one seems relatively nice and clear, I can try something like that 👍 - I'm still finding it a bit difficult to remember how to wrap my head around the sorting and filtering going on each time I look at it, so I'm inclined to accept a bit of overhead if it makes things easier to follow. I can try something similar for the undirected version too.
| for other_node_type in np.unique(other_node_types): | ||
|
|
||
| # filter node ilocs based on other node type | ||
| arr_filtered = arr[other_node_types == other_node_type] | ||
| splits = _get_splits( | ||
| [arr_filtered], self.number_of_nodes, self._id_index.dtype | ||
| ) | ||
|
|
||
| # filter edge ilocs based on other node type | ||
| flat_filtered = flat[other_node_types[flat] == other_node_type] |
There was a problem hiding this comment.
The double indexing here is potentially suboptimal, especially because other_node_types[flat] doesn't change between loops but does have to be recomputed. I wonder if there's another way we could phrase this loop to be slightly more efficient even if we don't completely switch it:
other_node_types = other_node_types[flat]
for other_node_type in np.unique(other_node_types):
# filter edge ilocs based on other node type
flat_filtered = flat[other_node_types == other_node_type]
# choose the relevant node ilocs based on the edge ilocs
arr_filtered = arr[flat_filtered]
splits = _get_splits(
[arr_filtered], self.number_of_nodes, self._id_index.dtype
)(Not a GitHub suggestion because the rendering is really ugly.)
This reduces us from doing two O(edges) comparisons to doing one O(edges) one, and one O(selected edges) one. This doesn't change the overall asymptotics of the loop but should reduce the overhead. (Moving the other_node_types[flat] out of the loop could be applied even without changing the order of flat_filtered and arr_filtered.)
| self._create_undirected_adj_lists_by_other_node_type() | ||
| ) | ||
|
|
||
| def _create_undirected_adj_lists_by_other_node_type(self): |
There was a problem hiding this comment.
This function is... impressive. It definitely seems super hard to phrase some of this.
| return self._adj_lookup_by_other_node_type(ins=ins, outs=outs).get( | ||
| other_node_type_iloc, defaultdict(list) | ||
| )[node_iloc] |
There was a problem hiding this comment.
I take it the get+defaultdict(list) here is in case there's node types that don't have any incoming edges or any outgoing edges (e.g. all edges might be a->b)? I'm slightly concerned that creating a defaultdict for every call may be adding noticeable overhead, and, also it's returning a list rather than a NumPy array, so consumers may get confused (and/or have to spend more work converting types).
One option might be pre-filling the _edges..._dict_by_other_node_type values with every node type, e.g.:
class EmptyFlatAdjacencyList:
def __init__(self, dtype):
self.empty = np.array([], dtype=dtype)
def __getitem__(self, idx):
return emptyOr just creating an O(nodes) FlatAdjacencyList, like FlatAdjacencyList(np.array([], dtype=...), np.zeros(len(nodes) + 1))
Another option would be moving the defaulting to a separate code path, like:
adj_list = self._adj_lookup_by_other_node_type(ins=ins, outs=outs).get(other_node_type_iloc)
if adj_list is None:
return np.array([], dtype=...)
return adj_list[node_iloc]There was a problem hiding this comment.
Ah good point, I did want to also allow completely non-existent node types to return empty lists:
g = StellarGraph(...) # has node types ["A", "B"]
g.neighbors(1, other_node_type="X") # should return np.array([])So I might go with the separate code path approach?
There was a problem hiding this comment.
What's the motivation for returning empty, rather than emitting an error? Because this is using ilocs, I suspect the conversion to a type iloc should fail for that case?
| @@ -1001,14 +1005,30 @@ def f(): | |||
| @pytest.mark.parametrize( | |||
| "num_nodes,num_edges", [(100, 200), (1000, 5000), (20000, 100000)] | |||
| ) | |||
| @pytest.mark.parametrize("force_adj_lists", ["directed", "undirected"]) | |||
| @pytest.mark.parametrize( | |||
There was a problem hiding this comment.
In order to capture some of performance of the filtering in the inner loop, I wonder if if it's worth also parameterising this benchmark by the number of node types, e.g.
@pytest.mark.parametrize("n_types", [1, 4, 20])and passing this through to example_benchmark_graph.
(Potentially with a if force_adj_lists in ("directed", "undirected") and n_types != 1: pytest.skip("...") to avoid spending time on retesting uninteresting things.)
|
Is this enhancement already in place? How can I start using this version? The current implementation of |
Pulled out while working on #934 - currently this demo is a bit unusable and doesn't provide a good enough experience to be converted into a demo notebook, since it takes hours for it to run on a laptop, but with the changes from this PR, it should run within ~ 10 minutes or so depending on the parameters used.
This mainly updates
EdgeDatato store an inner set of dicts within each adjacency list, so that we group adjacency elements by their node type. This could probably be useful for other heterogeneous algorithms that want to traverse the graph either for random walks or SAGE-type sampling - currently it's just been used for Metapath2Vec specifically. At its current state in this PR, the speedup for Metapath2vec looks pretty great, butgraph.neighboursbecomes quite a bit slower for homogeneous graphsBenchmark against
developfor metapath walk:Things do get a bit slower for just getting neighbours without caring about types, and constructing the stellargraph object