Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
torch-cluster
Commits
3aad518a
Commit
3aad518a
authored
Mar 25, 2018
by
rusty1s
Browse files
added normalized cut implementation
parent
af96c80a
Changes
3
Hide whitespace changes
Inline
Side-by-side
Showing
3 changed files
with
44 additions
and
4 deletions
+44
-4
test/test_graclus.py
test/test_graclus.py
+2
-2
torch_cluster/functions/degree.py
torch_cluster/functions/degree.py
+7
-0
torch_cluster/functions/graclus.py
torch_cluster/functions/graclus.py
+35
-2
No files found.
test/test_graclus.py
View file @
3aad518a
import
pytest
import
torch
import
torch
from
torch_cluster
import
graclus_cluster
from
torch_cluster
import
graclus_cluster
...
@@ -7,5 +6,6 @@ def test_graclus():
...
@@ -7,5 +6,6 @@ def test_graclus():
edge_index
=
torch
.
LongTensor
([[
0
,
0
,
0
,
1
,
2
,
3
,
3
,
3
,
4
,
5
,
5
,
5
,
6
,
6
],
edge_index
=
torch
.
LongTensor
([[
0
,
0
,
0
,
1
,
2
,
3
,
3
,
3
,
4
,
5
,
5
,
5
,
6
,
6
],
[
2
,
3
,
6
,
5
,
0
,
0
,
4
,
5
,
3
,
1
,
3
,
6
,
0
,
3
]])
[
2
,
3
,
6
,
5
,
0
,
0
,
4
,
5
,
3
,
1
,
3
,
6
,
0
,
3
]])
edge_attr
=
torch
.
Tensor
([
2
,
2
,
2
,
1
,
2
,
2
,
2
,
2
,
2
,
1
,
2
,
2
,
2
,
2
])
edge_attr
=
torch
.
Tensor
([
2
,
2
,
2
,
1
,
2
,
2
,
2
,
2
,
2
,
1
,
2
,
2
,
2
,
2
])
rid
=
torch
.
LongTensor
([
0
,
1
,
2
,
3
,
4
,
5
,
6
])
graclus_cluster
(
edge_index
,
edge_attr
)
graclus_cluster
(
edge_index
,
edge_attr
=
edge_attr
,
rid
=
rid
)
torch_cluster/functions/degree.py
0 → 100644
View file @
3aad518a
import
torch
def
node_degree
(
edge_index
,
num_nodes
,
out
=
None
):
zero
=
torch
.
zeros
(
num_nodes
,
out
=
out
)
one
=
torch
.
ones
(
edge_index
.
size
(
1
),
out
=
zero
.
new
())
return
zero
.
scatter_add_
(
0
,
edge_index
[
0
],
one
)
torch_cluster/functions/graclus.py
View file @
3aad518a
def
graclus_cluster
(
edge_index
,
edge_attr
=
None
,
batch
=
None
):
from
__future__
import
division
pass
import
torch
from
.degree
import
node_degree
def
graclus_cluster
(
edge_index
,
num_nodes
=
None
,
edge_attr
=
None
,
batch
=
None
,
rid
=
None
):
num_nodes
=
edge_index
.
max
()
+
1
if
num_nodes
is
None
else
num_nodes
rid
=
torch
.
randperm
(
num_nodes
)
if
rid
is
None
else
rid
cut
=
normalized_cut
(
edge_index
,
num_nodes
,
edge_attr
)
print
(
cut
)
def
normalized_cut
(
edge_index
,
num_nodes
,
edge_attr
=
None
):
row
,
col
=
edge_index
out
=
edge_attr
.
new
()
if
edge_attr
is
not
None
else
torch
.
Tensor
()
cut
=
node_degree
(
edge_index
,
num_nodes
,
out
=
out
)
cut
=
1
/
cut
cut
=
cut
[
row
]
+
cut
[
col
]
if
edge_attr
is
None
:
return
cut
else
:
if
edge_attr
.
dim
()
>
1
and
edge_attr
.
size
(
1
)
>
1
:
edge_attr
=
torch
.
norm
(
edge_attr
,
2
,
1
)
return
edge_attr
.
squeeze
()
*
cut
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment