Merge pull request #69 from yuehhua/develop

yuehhua · web-flow · commit 7e19a7bcdfd7 · 2020-07-20T11:06:46.000+08:00
Graph API enhancements
diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
@@ -1,18 +1,36 @@
 variables:
-  CI_IMAGE_TAG: 'cuda'
+  CI_APT_INSTALL: 'libgomp1'
   JULIA_NUM_THREADS: '4'
+  NVIDIA_VISIBLE_DEVICES: 'all'
+  NVIDIA_DRIVER_CAPABILITIES: 'compute,utility'
+  CI_THOROUGH: 'true'
+  JULIA_CUDA_VERSION: '10.2'
+  JULIA_CUDA_USE_BINARYBUILDER: 'true'
 
 include:
   - 'https://raw.githubusercontent.com/JuliaGPU/gitlab-ci/master/templates/v6.yml'
 
-image: nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04
+image: ubuntu:bionic
 
 test:v1.4:
   extends:
     - .julia:1.4
     - .test
   variables:
     CI_VERSION_TAG: 'v1.4'
+  tags:
+    - nvidia
+    - cuda_10.2
+
+test:v1.5:
+  extends:
+    - .julia:1.5
+    - .test
+  variables:
+    CI_VERSION_TAG: 'v1.5'
+  tags:
+    - nvidia
+    - cuda_10.2
 
 # test:dev:
 #   extends:
diff --git a/src/GeometricFlux.jl b/src/GeometricFlux.jl
@@ -85,7 +85,7 @@ export
     pool,
 
     # graph/index
-    neighbors,
+    adjacency_list,
     generate_cluster,
 
     # graph/featuredgraphs
@@ -102,9 +102,6 @@ export
     has_global_feature,
     nv,
 
-    # graph/simplegraphs
-    adjlist,
-
     # utils
     gather,
     topk_index
@@ -115,12 +112,12 @@ include("operations/scatter.jl")
 include("operations/pool.jl")
 include("operations/linalg.jl")
 
-include("utils.jl")
-
 include("graph/index.jl")
 include("graph/featuredgraphs.jl")
 include("graph/linalg.jl")
 
+include("utils.jl")
+
 include("layers/gn.jl")
 include("layers/msgpass.jl")
 
diff --git a/src/graph/featuredgraphs.jl b/src/graph/featuredgraphs.jl
@@ -81,12 +81,12 @@ has_global_feature(::NullGraph) = false
 has_global_feature(fg::FeaturedGraph) = fg.gf[] != zeros(0)
 
 """
-    neighbors(::AbstractFeaturedGraph)
+    adjacency_list(::AbstractFeaturedGraph)
 
 Get adjacency list of graph.
 """
-neighbors(::NullGraph) = [zeros(0)]
-neighbors(fg::FeaturedGraph) = neighbors(fg.graph[])
+adjacency_list(::NullGraph) = [zeros(0)]
+adjacency_list(fg::FeaturedGraph) = adjacency_list(fg.graph[])
 
 """
     nv(::AbstractFeaturedGraph)
diff --git a/src/graph/index.jl b/src/graph/index.jl
@@ -1,9 +1,9 @@
 """
-    neighbors(adj)
+    adjacency_list(adj)
 
 Transform a adjacency matrix into a adjacency list.
 """
-function neighbors(adj::AbstractMatrix{T}) where {T}
+function adjacency_list(adj::AbstractMatrix{T}) where {T}
     n = size(adj,1)
     @assert n == size(adj,2) "adjacency matrix is not a square matrix."
     A = (adj .!= zero(T))
@@ -15,7 +15,9 @@ function neighbors(adj::AbstractMatrix{T}) where {T}
     return ne
 end
 
-neighbors(adj::AbstractVector{<:AbstractVector{<:Integer}}) = adj
+adjacency_list(adj::AbstractVector{<:AbstractVector{<:Integer}}) = adj
+
+Zygote.@nograd adjacency_list
 
 """
     accumulated_edges(adj[, num_V])
@@ -30,6 +32,8 @@ function accumulated_edges(adj::AbstractVector{<:AbstractVector{<:Integer}},
     y
 end
 
+Zygote.@nograd accumulated_edges
+
 Zygote.@nograd function generate_cluster(M::AbstractArray{T,N}, accu_edge, V, E) where {T,N}
     cluster = similar(M, Int, E)
     @inbounds for i = 1:V
@@ -68,6 +72,8 @@ function vertex_pair_table(eidx::Dict)
     table
 end
 
+Zygote.@nograd vertex_pair_table
+
 """
     edge_index_table(adj[, num_E])
 
@@ -96,6 +102,8 @@ function edge_index_table(vpair::AbstractVector{<:Tuple})
     table
 end
 
+Zygote.@nograd edge_index_table
+
 function transform(X::AbstractArray, vpair::AbstractVector{<:Tuple}, num_V)
     dims = size(X)[1:end-1]..., num_V, num_V
     Y = similar(X, dims)
diff --git a/src/graph/simplegraphs.jl b/src/graph/simplegraphs.jl
@@ -1,6 +1,11 @@
-using LightGraphs: AbstractSimpleGraph, nv, adjacency_matrix, inneighbors,
+using LightGraphs: AbstractSimpleGraph, nv, adjacency_matrix, inneighbors, outneighbors,
                    all_neighbors
 
+function adjacency_list(g::AbstractSimpleGraph)
+    N = nv(g)
+    Vector{Int}[outneighbors(g, i) for i = 1:N]
+end
+
 ## Convolution layers accepting AbstractSimpleGraph
 
 function GCNConv(g::AbstractSimpleGraph, ch::Pair{<:Integer,<:Integer}, σ = identity;
diff --git a/src/graph/weightedgraphs.jl b/src/graph/weightedgraphs.jl
@@ -1,4 +1,9 @@
-using SimpleWeightedGraphs: AbstractSimpleWeightedGraph, nv
+using SimpleWeightedGraphs: AbstractSimpleWeightedGraph, nv, outneighbors
+
+function adjacency_list(g::AbstractSimpleWeightedGraph)
+    N = nv(g)
+    Vector{Int}[outneighbors(g, i) for i = 1:N]
+end
 
 ## Convolution layers accepting AbstractSimpleWeightedGraph
 
diff --git a/src/layers/conv.jl b/src/layers/conv.jl
@@ -198,7 +198,7 @@ function GraphConv(adj::AbstractMatrix, ch::Pair{<:Integer,<:Integer}, aggr=:add
     w1 = T.(init(ch[2], ch[1]))
     w2 = T.(init(ch[2], ch[1]))
     b = bias ? T.(init(ch[2])) : zeros(T, ch[2])
-    fg = FeaturedGraph(neighbors(adj))
+    fg = FeaturedGraph(adjacency_list(adj))
     GraphConv(fg, w1, w2, b, aggr)
 end
 
@@ -264,7 +264,7 @@ function GATConv(adj::AbstractMatrix, ch::Pair{<:Integer,<:Integer}; heads::Inte
     w = T.(init(ch[2]*heads, ch[1]))
     b = bias ? T.(init(ch[2]*heads)) : zeros(T, ch[2]*heads)
     a = T.(init(2*ch[2], heads, 1))
-    fg = FeaturedGraph(neighbors(adj))
+    fg = FeaturedGraph(adjacency_list(adj))
     GATConv(fg, w, b, a, negative_slope, ch, heads, concat)
 end
 
@@ -349,7 +349,7 @@ function GatedGraphConv(adj::AbstractMatrix, out_ch::Integer, num_layers::Intege
                         aggr=:add, init=glorot_uniform, T::DataType=Float32)
     w = T.(init(out_ch, out_ch, num_layers))
     gru = GRUCell(out_ch, out_ch)
-    fg = FeaturedGraph(neighbors(adj))
+    fg = FeaturedGraph(adjacency_list(adj))
     GatedGraphConv(fg, w, gru, out_ch, num_layers, aggr)
 end
 
@@ -415,7 +415,7 @@ struct EdgeConv{V<:AbstractFeaturedGraph} <: MessagePassing
 end
 
 function EdgeConv(adj::AbstractMatrix, nn; aggr::Symbol=:max)
-    fg = FeaturedGraph(neighbors(adj))
+    fg = FeaturedGraph(adjacency_list(adj))
     EdgeConv(fg, nn, aggr)
 end
 
diff --git a/src/layers/gn.jl b/src/layers/gn.jl
@@ -62,7 +62,7 @@ end
 end
 
 function propagate(gn::T, fg::FeaturedGraph, naggr=nothing, eaggr=nothing, vaggr=nothing) where {T<:GraphNet}
-    adj = neighbors(fg)
+    adj = adjacency_list(fg)
     num_V = nv(fg)
     accu_edge = accumulated_edges(adj)
     num_E = accu_edge[end]
diff --git a/src/layers/msgpass.jl b/src/layers/msgpass.jl
@@ -56,7 +56,7 @@ end
 end
 
 function propagate(mp::T, fg::FeaturedGraph, aggr::Symbol=:add) where {T<:MessagePassing}
-    adj = neighbors(fg)
+    adj = adjacency_list(fg)
     num_V = nv(fg)
     accu_edge = accumulated_edges(adj)
     num_E = accu_edge[end]
diff --git a/src/operations/linalg.jl b/src/operations/linalg.jl
@@ -2,7 +2,7 @@
 
 Zygote.@nograd issymmetric
 
-function adjacency_matrix(adj::AbstractMatrix, T::DataType=eltype(adj))
+Zygote.@nograd function adjacency_matrix(adj::AbstractMatrix, T::DataType=eltype(adj))
     m, n = size(adj)
     (m == n) || throw(DimensionMismatch("adjacency matrix is not a square matrix: ($m, $n)"))
     T.(adj)
@@ -33,6 +33,11 @@ julia> GeometricFlux.degrees(m)
 ```
 """
 function degrees(adj::AbstractMatrix, T::DataType=eltype(adj); dir::Symbol=:out)
+    _degrees(T.(adj), dir)
+end
+
+# nograd can only used without keyword arguments
+Zygote.@nograd function _degrees(adj::AbstractMatrix, dir::Symbol=:out)
     if issymmetric(adj)
         d = vec(sum(adj, dims=1))
     else
@@ -90,7 +95,7 @@ The values other than diagonal are zeros.
 - `dir`: direction of degree; should be `:in`, `:out`, or `:both` (optional).
 """
 function inv_sqrt_degree_matrix(adj::AbstractMatrix, T::DataType=eltype(adj); dir::Symbol=:out)
-    d  = inv.(sqrt.(degrees(adj, T, dir=dir)))
+    d = inv.(sqrt.(degrees(adj, T, dir=dir)))
     return Diagonal(d)
 end
 
diff --git a/src/utils.jl b/src/utils.jl
@@ -81,3 +81,11 @@ topk_index(y::Adjoint, k::Integer) = topk_index(y', k)
 
 get_feature(::Nothing, i::Integer) = zeros(0)
 get_feature(A::AbstractMatrix, i::Integer) = (i ≤ size(A,2)) ? view(A, :, i) : zeros(0)
+
+
+function bypass_graph(nf_func=identity, ef_func=identity, gf_func=identity)
+    return function (fg::FeaturedGraph)
+        FeaturedGraph(graph(fg), nf_func(node_feature(fg)), ef_func(edge_feature(fg)),
+                      gf_func(global_feature(fg)))
+    end
+end
diff --git a/test/operations/linalg.jl b/test/operations/linalg.jl
@@ -33,11 +33,11 @@
         for T in [Float16, Float32, Float64]
             @test normalized_laplacian(adj, T) ≈ T.(norm_lap)
             @test eltype(normalized_laplacian(adj, T)) == T
-            
+
             @test scaled_laplacian(adj, T) ≈ T.(scaled_lap)
             @test eltype(scaled_laplacian(adj, T)) == T
         end
-        @test neighbors(adj) == [[2,4], [1,3], [2,4], [1,3]]
+        @test adjacency_list(adj) == [[2,4], [1,3], [2,4], [1,3]]
     end
 
     @testset "asymmetric" begin
@@ -74,6 +74,6 @@
             @test eltype(laplacian_matrix(adj, T, dir=:in)) == T
             @test eltype(laplacian_matrix(adj, T, dir=:both)) == T
         end
-        @test neighbors(adj) == [[2,3,4], [1,3], [1,2,4], [1,3]]
+        @test adjacency_list(adj) == [[2,3,4], [1,3], [1,2,4], [1,3]]
     end
 end
