HeteroGraphConv documentation (#312)

* batching

* HeteroGraphConv docs

* complete docs

* more tests

* fix test

Author: CarloLucibello

docs/src/api/heterograph.md

@@ -1,11 +1,21 @@
-# GNNHeteroGraph
+# Hetereogeneous Graphs
 
-Documentation page for the graph type `GNNHeteroGraph` and related methods representing heterogeneous graphs,
-where  nodes and edges can have different types.
+
+## GNNHeteroGraph
+Documentation page for the type `GNNHeteroGraph` representing heterogeneous graphs, where  nodes and edges can have different types.
 
 
 ```@autodocs
 Modules = [GraphNeuralNetworks.GNNGraphs]
 Pages   = ["gnnheterograph.jl"]
 Private = false
 ```
+
+## Heterogeneous Graph Convolutions
+
+Heterogeneous graph convolutions are implemented in the type [`HeteroGraphConv`](@ref).
+`HeteroGraphConv` relies on standard graph convolutional layers to perform message passing on the different relations. See the table at [this page](https://carlolucibello.github.io/GraphNeuralNetworks.jl/dev/api/conv/) for the supported layers.
+
+```@docs
+HeteroGraphConv
+```

docs/src/gnngraph.md

@@ -189,15 +189,15 @@ graph_indicator(gall)
 ## DataLoader and mini-batch iteration
 
 While constructing a batched graph and passing it to the `DataLoader` is always 
-an option for mini-batch iteration, the recommended way is
-to pass an array of graphs directly:
+an option for mini-batch iteration, the recommended way for better performance is
+to pass an array of graphs directly and set the `collate` option to `true`:
 
 ```julia
 using Flux: DataLoader
 
 data = [rand_graph(10, 30, ndata=rand(Float32, 3, 10)) for _ in 1:320]
 
-train_loader = DataLoader(data, batchsize=16, shuffle=true)
+train_loader = DataLoader(data, batchsize=16, shuffle=true, collate=true)
 
 for g in train_loader
     @assert g.num_graphs == 16

docs/src/heterograph.md

@@ -94,4 +94,37 @@ GNNHeteroGraph:
         (:user, :rate, :movie)  =>  DataStore(e = [64×4 Matrix{Float32}])
 ```
 
+## Batching
+Similarly to graphs, also heterographs can be batched together.
+```julia-repl
+```julia
+julia> gs = [rand_bipartite_heterograph((5, 10), 20) for _ in 1:32];
+
+julia> Flux.batch(gs)
+GNNHeteroGraph:
+  num_nodes: Dict(:A => 160, :B => 320)
+  num_edges: Dict((:A, :to, :B) => 640, (:B, :to, :A) => 640)
+  num_graphs: 32
+```
+Batching is automatically performed by the [`DataLoader`](@ref) iterator
+when the `collate` option is set to `true`.
+```julia-repl
+using Flux: DataLoader
+
+data = [rand_bipartite_heterograph((5, 10), 20, 
+            ndata=Dict(:A=>rand(Float32, 3, 5))) 
+        for _ in 1:320];
+
+train_loader = DataLoader(data, batchsize=16, shuffle=true, collate=true)
+
+for g in train_loader
+    @assert g.num_graphs == 16
+    @assert g.num_nodes[:A] == 80
+    @assert size(g.ndata[:A].x) == (3, 80)    
+    # ...
+end
+```
+
+## Graph convolutions on heterographs
 
+See [`HeteroGraphConv`](@ref) for how to perform convolutions on heterogenous graphs.

src/layers/heteroconv.jl

@@ -1,3 +1,41 @@
+@doc raw"""
+    HeteroGraphConv(itr; aggr = +)
+    HeteroGraphConv(pairs...; aggr = +)
+
+A convolutional layer for heterogeneous graphs.
+
+The `itr` argument is an iterator of `pairs` of the form `edge_t => layer`, where `edge_t` is a
+3-tuple of the form `(src_node_type, edge_type, dst_node_type)`, and `layer` is a 
+convolutional layers for homogeneous graphs. 
+
+Each convolution is applied to the corresponding relation. 
+Since a node type can be involved in multiple relations, the single convolution outputs 
+have to be aggregated using the `aggr` function. The default is to sum the outputs.
+
+# Forward Arguments
+
+* `g::GNNHeteroGraph`: The input graph.
+* `x::Union{NamedTuple,Dict}`: The input node features. The keys are node types and the
+  values are node feature tensors.
+
+# Examples 
+
+```julia-repl
+julia> g = rand_bipartite_heterograph((10, 15), 20)
+GNNHeteroGraph:
+  num_nodes: Dict(:A => 10, :B => 15)
+  num_edges: Dict((:A, :to, :B) => 20, (:B, :to, :A) => 20)
+
+julia> x = (A = rand(Float32, 64, 10), B = rand(Float32, 64, 15));
+
+julia> layer = HeteroGraphConv((:A, :to, :B) => GraphConv(64 => 32, relu),
+                               (:B, :to, :A) => GraphConv(64 => 32, relu));
+
+julia> y = layer(g, x); # output is a named tuple
+
+julia> size(y.A) == (32, 10) && size(y.B) == (32, 15)
+true
+"""
 struct HeteroGraphConv
     etypes::Vector{EType}
     layers::Vector{<:GNNLayer}
@@ -6,6 +44,9 @@ end
 
 Flux.@functor HeteroGraphConv
 
+HeteroGraphConv(itr::Dict; aggr = +) = HeteroGraphConv(pairs(itr); aggr)
+HeteroGraphConv(itr::Pair...; aggr = +) = HeteroGraphConv(itr; aggr)
+
 function HeteroGraphConv(itr; aggr = +)
     etypes = [k[1] for k in itr]
     layers = [k[2] for k in itr]
@@ -37,3 +78,17 @@ function _reduceby_node_t(aggr, outs, ntypes)
     vals = [_reduce(node_t) for node_t in ntypes]
     return NamedTuple{tuple(ntypes...)}(vals)
 end
+
+function Base.show(io::IO, hgc::HeteroGraphConv)
+    if get(io, :compact, false)
+        print(io, "HeteroGraphConv(aggr=$(hgc.aggr))")
+    else
+        println(io, "HeteroGraphConv(aggr=$(hgc.aggr)):")
+        for (i, (et,layer)) in enumerate(zip(hgc.etypes, hgc.layers))
+            print(io, "  $(et => layer)")
+            if i < length(hgc.etypes)
+                print(io, "\n")
+            end
+        end
+    end
+end

test/layers/heteroconv.jl

@@ -27,4 +27,10 @@
         @test dx[:A] ≈ ndx[:A]        rtol=1e-4
         @test dx[:B] ≈ ndx[:B]        rtol=1e-4
     end
+
+    @testset "Constructor from pairs" begin
+        layer = HeteroGraphConv((:A, :to, :B) => GraphConv(64 => 32, relu),
+                                (:B, :to, :A) => GraphConv(64 => 32, relu));
+        @test length(layer.etypes) == 2
+    end
 end