add DeepSet model

fix

Author: yuehhua

docs/src/manual/models.md

@@ -37,6 +37,22 @@ Reference: [Variational Graph Auto-Encoders](https://arxiv.org/abs/1611.07308)
 
 ---
 
+## DeepSet
+
+```math
+Z = \rho ( \sum_{x_i \in \mathcal{V}} \phi (x_i) )
+```
+
+where ``\phi`` and ``\rho`` denote two neural networks and ``x_i`` is the node feature for node ``i``.
+
+```@docs
+GeometricFlux.DeepSet
+```
+
+Reference: [Deep Sets](https://papers.nips.cc/paper/2017/hash/f22e4747da1aa27e363d86d40ff442fe-Abstract.html)
+
+---
+
 ## Special Layers
 
 ### Inner-product Decoder

examples/deepsets.jl

@@ -0,0 +1,3 @@
+mae loss
+
+Adam(lr=1e-4)

src/GeometricFlux.jl

@@ -51,6 +51,7 @@ export
     VGAE,
     InnerProductDecoder,
     VariationalGraphEncoder,
+    DeepSet,
 
     # layer/utils
     WithGraph,

src/models.jl

@@ -146,3 +146,77 @@ WithGraph(fg::AbstractFeaturedGraph, l::VariationalGraphEncoder) =
         WithGraph(fg, l.logσ),
         l.z_dim
     )
+
+
+"""
+    DeepSet(ϕ, ρ, aggr=+)
+
+Deep set model.
+
+# Arguments
+
+- `ϕ`: The dimension of input features.
+- `ρ`: The dimension of output features.
+- `aggr`: An aggregate function applied to the result of message function. `+`, `-`,
+`*`, `/`, `max`, `min` and `mean` are available.
+
+# Examples
+
+```jldoctest
+julia> ϕ = Dense(64, 16)
+Dense(64, 16)       # 1_040 parameters
+
+julia> ρ = Dense(16, 4)
+Dense(16, 4)        # 68 parameters
+
+julia> DeepSet(ϕ, ρ)
+DeepSet(Dense(64, 16), Dense(16, 4), aggr=+)
+
+julia> DeepSet(ϕ, ρ, aggr=max)
+DeepSet(Dense(64, 16), Dense(16, 4), aggr=max)
+```
+
+See also [`WithGraph`](@ref) for training layer with static graph.
+"""
+struct DeepSet{T,S,O} <: GraphNet
+    ϕ::T
+    ρ::S
+    aggr::O
+end
+
+DeepSet(ϕ, ρ; aggr=+) = DeepSet(ϕ, ρ, aggr)
+
+update_vertex(l::DeepSet, Ē, V, u) = l.ϕ(V)
+
+update_global(l::DeepSet, ē, v̄, u) = l.ρ(v̄)
+
+# For variable graph
+function (l::DeepSet)(fg::AbstractFeaturedGraph)
+    X = node_feature(fg)
+    u = global_feature(fg)
+    GraphSignals.check_num_nodes(fg, X)
+    _, _, u = propagate(l, graph(fg), nothing, X, u, nothing, nothing, l.aggr)
+    return ConcreteFeaturedGraph(fg, gf=u)
+end
+
+# For static graph
+function (l::DeepSet)(el::NamedTuple, X::AbstractArray, u=nothing)
+    GraphSignals.check_num_nodes(el.N, X)
+    _, _, u = propagate(l, el, nothing, X, u, nothing, nothing, l.aggr)
+    return u
+end
+
+WithGraph(fg::AbstractFeaturedGraph, l::DeepSet) = WithGraph(to_namedtuple(fg), l)
+(wg::WithGraph{<:DeepSet})(args...) = wg.layer(wg.graph, args...)
+
+function Base.show(io::IO, l::DeepSet)
+    print(io, "DeepSet(", l.ϕ, ", ", l.ρ)
+    print(io, ", aggr=", l.aggr)
+    print(io, ")")
+end
+
+function Base.show(io::IO, l::WithGraph{<:DeepSet})
+    print(io, "WithGraph(Graph(#V=", l.graph.N)
+    print(io, ", #E=", l.graph.E, "), ")
+    print(io, l.layer, ")")
+end

src/operation.jl

@@ -14,12 +14,12 @@ function batched_index(idx::AbstractVector, batch_size::Integer)
     return tuple.(idx, b)
 end
 
-aggregate(aggr::typeof(+), X) = vec(sum(X, dims=2))
-aggregate(aggr::typeof(-), X) = -vec(sum(X, dims=2))
-aggregate(aggr::typeof(*), X) = vec(prod(X, dims=2))
-aggregate(aggr::typeof(/), X) = 1 ./ vec(prod(X, dims=2))
-aggregate(aggr::typeof(max), X) = vec(maximum(X, dims=2))
-aggregate(aggr::typeof(min), X) = vec(minimum(X, dims=2))
-aggregate(aggr::typeof(mean), X) = vec(aggr(X, dims=2))
+aggregate(::typeof(+), X) = sum(X, dims=2)
+aggregate(::typeof(-), X) = -sum(X, dims=2)
+aggregate(::typeof(*), X) = prod(X, dims=2)
+aggregate(::typeof(/), X) = 1 ./ prod(X, dims=2)
+aggregate(::typeof(max), X) = maximum(X, dims=2)
+aggregate(::typeof(min), X) = minimum(X, dims=2)
+aggregate(::typeof(mean), X) = mean(X, dims=2)
 
 @non_differentiable batched_index(x...)

test/models.jl

@@ -1,4 +1,5 @@
 @testset "models" begin
+    batch_size = 10
     in_channel = 3
     out_channel = 5
     N = 4
@@ -60,5 +61,31 @@
             Y = node_feature(fg_)
             @test size(Y) == (N, N)
         end
+
+        @testset "DeepSet" begin
+            ϕ = Dense(64, 16)
+            ρ = Dense(16, 4)
+            @testset "layer without graph" begin
+                deepset = DeepSet(ϕ, ρ)
+
+                X = rand(T, 64, N)
+                fg = FeaturedGraph(adj, nf=X)
+                fg_ = deepset(fg)
+                @test size(global_feature(fg_)) == (4, 1)
+
+                g = Zygote.gradient(() -> sum(global_feature(deepset(fg))), Flux.params(deepset))
+                @test length(g.grads) == 2
+            end
+
+            @testset "layer with static graph" begin
+                X = rand(T, 64, N, batch_size)
+                deepset = WithGraph(fg, DeepSet(ϕ, ρ))
+                Y = deepset(X)
+                @test size(Y) == (4, 1, batch_size)
+
+                g = Zygote.gradient(() -> sum(deepset(X)), Flux.params(deepset))
+                @test length(g.grads) == 0
+            end
+        end
     end
 end