Skip to content

GitLab

"src/array/cuda/git@developer.sourcefind.cn:OpenDAS/dgl.git" did not exist on "7b9afbfa28e9441fd3bfbe4e3c24eb62816cb1a7"
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
Hide 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