Merge pull request #269 from FluxML/gae

Fix GAE example

Author: yuehhua

examples/gae.jl

@@ -1,45 +1,112 @@
-using GeometricFlux
-using GraphSignals
+using CUDA
 using Flux
-using Flux: throttle
+using Flux: onecold
 using Flux.Losses: logitbinarycrossentropy
-using Flux: @epochs
-using JLD2
-using Statistics: mean
-using SparseArrays
-using Graphs.SimpleGraphs
-using CUDA
+using Flux.Data: DataLoader
+using GeometricFlux
+using GeometricFlux.Datasets
+using GraphSignals
+using Parameters: @with_kw
+using ProgressMeter: Progress, next!
+using Statistics
+using Random
+
+function load_data(dataset, batch_size, train_repeats=128)
+    # (train_X, train_y) dim: (num_features, target_dim) × 1708
+    train_X, _ = map(x -> Matrix(x), alldata(Planetoid(), dataset))
+    # (test_X, test_y) dim: (num_features, target_dim) × 1000
+    test_X, _ = map(x -> Matrix(x), testdata(Planetoid(), dataset))
+    g = graphdata(Planetoid(), dataset)
+
+    X = hcat(train_X, test_X)
+    fg = FeaturedGraph(g)
+    A = GraphSignals.adjacency_matrix(fg)
+    data = (repeat(X, outer=(1,1,train_repeats)), repeat(A, outer=(1,1,train_repeats)))
+    loader = DataLoader(data, batchsize=batch_size, shuffle=true)
+    return loader, fg
+end
+
+@with_kw mutable struct Args
+    η = 0.01                # learning rate
+    batch_size = 16         # batch size
+    epochs = 200            # number of epochs
+    seed = 0                # random seed
+    cuda = true             # use GPU
+    input_dim = 1433        # input dimension
+    hidden1_dim = 32        # hidden1 dimension
+    hidden2_dim = 16        # hidden1 dimension
+end
+
+## Loss: binary cross entropy
+model_loss(model, X, A) = logitbinarycrossentropy(model(X), A)
+
+function precision(model, X::AbstractArray, A::AbstractArray)
+    ŷ = onecold(softmax(cpu(model(X))))
+    y = onecold(cpu(A))
+    return mean(y[ŷ .== true])
+end
+
+precision(model, loader::DataLoader, device) =
+    mean(precision(model, X |> device, A |> device) for (X, A) in loader)
+
+function train(; kws...)
+    # load hyperparamters
+    args = Args(; kws...)
+    args.seed > 0 && Random.seed!(args.seed)
+
+    # GPU config
+    if args.cuda && CUDA.has_cuda()
+        device = gpu
+        @info "Training on GPU"
+    else
+        device = cpu
+        @info "Training on CPU"
+    end
+
+    # load Cora from Planetoid dataset
+    loader, fg = load_data(:cora, args.batch_size)
+    
+    # build model
+    encoder = Chain(
+        WithGraph(fg, GCNConv(args.input_dim=>args.hidden1_dim, relu)),
+        Dropout(0.5),
+        WithGraph(fg, GCNConv(args.hidden1_dim=>args.hidden2_dim)),
+    )
 
-CUDA.allowscalar(false)
+    model = GAE(encoder, σ) |> device
 
-@load "data/cora_features.jld2" features
-@load "data/cora_graph.jld2" g
+    # ADAM optimizer
+    opt = ADAM(args.η)
+    
+    # parameters
+    ps = Flux.params(model)
 
-num_nodes = 2708
-num_features = 1433
-hidden1 = 32
-hidden2 = 16
-target_catg = 7
-epochs = 200
+    # training
+    train_steps = 0
+    @info "Start Training, total $(args.epochs) epochs"
+    for epoch = 1:args.epochs
+        @info "Epoch $(epoch)"
+        progress = Progress(length(loader))
 
-## Preprocessing data
-fg = FeaturedGraph(g)  # pass to gpu together in model layers
-train_X = Matrix{Float32}(features) |> gpu  # dim: num_features * num_nodes
-train_y = fg |> GraphSignals.adjacency_matrix |> gpu  # dim: num_nodes * num_nodes
+        for (X, A) in loader
+            loss, back = Flux.pullback(ps) do
+                model_loss(model, X |> device, A |> device)
+            end
+            prec = precision(model, loader, device)
+            grad = back(1f0)
+            Flux.Optimise.update!(opt, ps, grad)
 
-## Model
-encoder = Chain(GCNConv(fg, num_features=>hidden1, relu),
-                GCNConv(fg, hidden1=>hidden2))
-model = Chain(GAE(encoder, σ)) |> gpu;
-# do not show model architecture, showing CuSparseMatrix will trigger errors
+            # progress meter
+            next!(progress; showvalues=[
+                (:loss, loss),
+                (:precision, prec),
+            ])
 
-## Loss
-loss(x, y) = logitbinarycrossentropy(model(x), y)
+            train_steps += 1
+        end
+    end
 
-## Training
-ps = Flux.params(model)
-train_data = [(train_X, train_y)]
-opt = ADAM(0.01)
-evalcb() = @show(loss(train_X, train_y))
+    return model, args
+end
 
-@epochs epochs Flux.train!(loss, ps, train_data, opt, cb=throttle(evalcb, 10))
+model, args = train()

src/GeometricFlux.jl

@@ -50,7 +50,7 @@ export
     GAE,
     VGAE,
     InnerProductDecoder,
-    VariationalEncoder,
+    VariationalGraphEncoder,
 
     # layer/utils
     WithGraph,

src/models.jl

@@ -1,14 +1,15 @@
 """
-    GAE(enc[, σ])
+    GAE(enc, [σ=identity])
 
 Graph autoencoder.
 
 # Arguments
 - `enc`: encoder. It can be any graph convolutional layer.
+- `σ`: Activation function for decoder.
 
 Encoder is specified by user and decoder will be `InnerProductDecoder` layer.
 """
-struct GAE{T,S}
+struct GAE{T,S} <: AbstractGraphLayer
     encoder::T
     decoder::S
 end
@@ -17,11 +18,14 @@ GAE(enc, σ::Function=identity) = GAE(enc, InnerProductDecoder(σ))
 
 @functor GAE
 
-function (g::GAE)(X::AbstractMatrix)
-    Z = g.encoder(X)
-    A = g.decoder(Z)
-    A
-end
+# For variable graph
+(l::GAE)(fg::AbstractFeaturedGraph) = fg |> l.encoder |> l.decoder
+
+# For static graph
+WithGraph(fg::AbstractFeaturedGraph, l::GAE) = GAE(WithGraph(fg, l.encoder), l.decoder)
+
+(l::GAE)(X::AbstractMatrix) = X |> l.encoder |> l.decoder
+(l::GAE)(X::AbstractArray) = X |> l.encoder |> l.decoder
 
 
 """
@@ -34,22 +38,24 @@ Variational graph autoencoder.
 
 Encoder is specified by user and decoder will be `InnerProductDecoder` layer.
 """
-struct VGAE{T,S}
+struct VGAE{T,S} <: AbstractGraphLayer
     encoder::T
     decoder::S
 end
 
 function VGAE(enc, h_dim::Integer, z_dim::Integer, σ::Function=identity)
-    VGAE(VariationalEncoder(enc, h_dim, z_dim), InnerProductDecoder(σ))
+    VGAE(VariationalGraphEncoder(enc, h_dim, z_dim), InnerProductDecoder(σ))
 end
 
 @functor VGAE
 
-function (g::VGAE)(fg::FeaturedGraph)
-    fg_ = g.encoder(fg)
-    fg_ = g.decoder(fg_)
-    fg_
-end
+# For variable graph
+(l::VGAE)(fg::AbstractFeaturedGraph) = fg |> l.encoder |> l.decoder
+
+# For static graph
+WithGraph(fg::AbstractFeaturedGraph, l::VGAE) = VGAE(WithGraph(fg, l.encoder), l.decoder)
+
+(l::VGAE)(X::AbstractArray) = X |> l.encoder |> l.decoder
 
 
 """
@@ -60,13 +66,15 @@ Inner-product decoder layer.
 # Arguments
 - `σ`: activation function.
 """
-struct InnerProductDecoder
-    σ
+struct InnerProductDecoder{F}
+    σ::F
 end
 
 @functor InnerProductDecoder
 
 (i::InnerProductDecoder)(Z::AbstractMatrix)::AbstractMatrix = i.σ.(Z'*Z)
+(i::InnerProductDecoder)(Z::AbstractArray)::AbstractArray =
+    i.σ.(NNlib.batched_mul(NNlib.batched_transpose(Z), Z))
 
 function (i::InnerProductDecoder)(fg::FeaturedGraph)
     Z = node_feature(fg)
@@ -76,9 +84,9 @@ end
 
 
 """
-    VariationalEncoder(nn, h_dim, z_dim)
+    VariationalGraphEncoder(nn, h_dim, z_dim)
 
-Variational encoder layer.
+Variational graph encoder layer.
 
 # Arguments
 - `nn`: neural network. It can be any graph convolutional layer.
@@ -87,33 +95,44 @@ Variational encoder layer.
 
 Encoder can be any graph convolutional layer.
 """
-struct VariationalEncoder
-    nn
-    μ
-    logσ
-    z_dim::Integer
+struct VariationalGraphEncoder{L,M,S,T<:Integer} <: AbstractGraphLayer
+    nn::L
+    μ::M
+    logσ::S
+    z_dim::T
 end
 
-function VariationalEncoder(nn, h_dim::Integer, z_dim::Integer)
-    VariationalEncoder(nn,
+function VariationalGraphEncoder(nn, h_dim::Integer, z_dim::Integer)
+    VariationalGraphEncoder(nn,
                        GCNConv(h_dim=>z_dim),
                        GCNConv(h_dim=>z_dim),
                        z_dim)
 end
 
-@functor VariationalEncoder
+@functor VariationalGraphEncoder
 
-function (ve::VariationalEncoder)(fg::FeaturedGraph)::FeaturedGraph
+function (ve::VariationalGraphEncoder)(fg::FeaturedGraph)::FeaturedGraph
     μ, logσ = summarize(ve, fg)
     Z = sample(μ, logσ)
     FeaturedGraph(fg, nf=Z)
 end
 
-function summarize(ve::VariationalEncoder, fg::FeaturedGraph)
+function summarize(ve::VariationalGraphEncoder, fg::FeaturedGraph)
     fg_ = ve.nn(fg)
     fg_μ, fg_logσ = ve.μ(fg_), ve.logσ(fg_)
     node_feature(fg_μ), node_feature(fg_logσ)
 end
 
 sample(μ::AbstractArray{T}, logσ::AbstractArray{T}) where {T<:Real} =
     μ + exp.(logσ) .* randn(T, size(logσ))
+
+# For static graph
+WithGraph(fg::AbstractFeaturedGraph, l::VariationalGraphEncoder) =
+    VariationalGraphEncoder(
+        WithGraph(fg, l.nn),
+        WithGraph(fg, l.μ),
+        WithGraph(fg, l.logσ),
+        l.z_dim
+    )
+
+# (l::VariationalGraphEncoder)(X::AbstractArray) = X |> l.encoder |> l.decoder

test/models.jl

@@ -2,6 +2,7 @@
     in_channel = 3
     out_channel = 5
     N = 4
+    batch = 4
     T = Float32
     adj = T[0. 1. 0. 1.;
             1. 0. 1. 0.;
@@ -11,7 +12,7 @@
     fg = FeaturedGraph(adj)
 
     @testset "GAE" begin
-        gc = GCNConv(fg, in_channel=>out_channel)
+        gc = WithGraph(fg, GCNConv(in_channel=>out_channel))
         gae = GAE(gc)
         X = rand(T, in_channel, N)
         Y = gae(X)
@@ -21,9 +22,9 @@
     @testset "VGAE" begin
         @testset "InnerProductDecoder" begin
            ipd = InnerProductDecoder(identity)
-           X = rand(T, 1, N)
+           X = rand(T, 1, N, batch)
            Y = ipd(X)
-           @test size(Y) == (N, N)
+           @test size(Y) == (N, N, batch)
 
            X = rand(T, 1, N)
            fg = FeaturedGraph(adj, nf=X)
@@ -38,10 +39,10 @@
            @test size(Y) == (N, N)
         end
 
-        @testset "VariationalEncoder" begin
+        @testset "VariationalGraphEncoder" begin
             z_dim = 2
             gc = GCNConv(in_channel=>out_channel)
-            ve = VariationalEncoder(gc, out_channel, z_dim)
+            ve = VariationalGraphEncoder(gc, out_channel, z_dim)
             X = rand(T, in_channel, N)
             fg = FeaturedGraph(adj, nf=X)
             fg_ = ve(fg)

test/runtests.jl

@@ -27,7 +27,7 @@ tests = [
     "layers/utils",
     "sampling",
     "embedding/node2vec",
-    # "models",
+    "models",
 ]
 
 if CUDA.functional()