Skip to content

GitLab

Unverified Commit 6f603bbf authored by Lingfan Yu's avatar Lingfan Yu Committed by GitHub
Browse files

[BugFix] Fix performance bug in EdgeSoftmax (#452) 

* fix performance bug in EdgeSoftmax

* minor
parent 681e5217
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
python/dgl/nn/pytorch/softmax.py
View file @ 6f603bbf
...@@ -84,6 +84,8 @@ class EdgeSoftmax(nn.Module): ...@@ -84,6 +84,8 @@ class EdgeSoftmax(nn.Module):
graph.apply_edges( graph.apply_edges(
lambda edges: {self._logits_name : th.exp(edges.data[self._logits_name] - lambda edges: {self._logits_name : th.exp(edges.data[self._logits_name] -
edges.dst[self._max_logits_name])}) edges.dst[self._max_logits_name])})
# pop out temporary feature _max_logits, otherwise get_ndata_name could have huge overhead
graph.ndata.pop(self._max_logits_name)
# compute normalizer # compute normalizer
graph.update_all(fn.copy_edge(self._logits_name, self._logits_name), graph.update_all(fn.copy_edge(self._logits_name, self._logits_name),
fn.sum(self._logits_name, self._normalizer_name)) fn.sum(self._logits_name, self._normalizer_name))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment