Merge pull request #334 from FluxML/new_features

Replace GraphMLDatasets in favor of MLDatasets and update docs

Author: yuehhua

Project.toml

@@ -10,10 +10,10 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
-GraphMLDatasets = "21828b05-d3b3-40ad-870e-a4bc2f52d5e8"
 GraphSignals = "3ebe565e-a4b5-49c6-aed2-300248c3a9c1"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
@@ -26,17 +26,17 @@ Word2Vec = "c64b6f0f-98cd-51d1-af78-58ae84944834"
 
 [compat]
 CUDA = "3"
-ChainRulesCore = "1.7"
+ChainRulesCore = "1"
 DataStructures = "0.18"
 FillArrays = "0.13"
 Flux = "0.12 - 0.13"
-GraphMLDatasets = "0.1"
 GraphSignals = "0.7"
 Graphs = "1"
+MLDatasets = "0.7"
 NNlib = "0.8"
 NNlibCUDA = "0.2"
 Optimisers = "0.2"
-Reexport = "1.1"
+Reexport = "1"
 StatsBase = "0.33"
 Word2Vec = "0.5"
 julia = "1.6"

data/cora_features.jld2

@@ -1,7 +1,9 @@
 [deps]
+DemoCards = "311a05b2-6137-4a5a-b473-18580a3d38b5"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+GeometricFlux = "7e08b658-56d3-11e9-2997-919d5b31e4ea"
 
 [compat]
 Documenter = "0.27"

data/cora_graph.jld2

@@ -201,6 +201,7 @@ @inproceedings{Satorras2021
 @article{Dwivedi2021,
    abstract = {Graph neural networks (GNNs) have become the standard learning architectures for graphs. GNNs have been applied to numerous domains ranging from quantum chemistry, recommender systems to knowledge graphs and natural language processing. A major issue with arbitrary graphs is the absence of canonical positional information of nodes, which decreases the representation power of GNNs to distinguish e.g. isomorphic nodes and other graph symmetries. An approach to tackle this issue is to introduce Positional Encoding (PE) of nodes, and inject it into the input layer, like in Transformers. Possible graph PE are Laplacian eigenvectors. In this work, we propose to decouple structural and positional representations to make easy for the network to learn these two essential properties. We introduce a novel generic architecture which we call LSPE (Learnable Structural and Positional Encodings). We investigate several sparse and fully-connected (Transformer-like) GNNs, and observe a performance increase for molecular datasets, from 2.87% up to 64.14% when considering learnable PE for both GNN classes.},
    author = {Vijay Prakash Dwivedi and Anh Tuan Luu and Thomas Laurent and Yoshua Bengio and Xavier Bresson},
+   journal = {ArXiv},
    month = {10},
    title = {Graph Neural Networks with Learnable Structural and Positional Representations},
    url = {http://arxiv.org/abs/2110.07875},

data/cora_label2onehot.jld2

@@ -1,30 +1,31 @@
 using Documenter
 using DocumenterCitations
+using DemoCards
 using GeometricFlux
 
+const ASSETS = ["assets/flux.css", "assets/favicon.ico"]
+
 bib = CitationBibliography(joinpath(@__DIR__, "bibliography.bib"), sorting=:nyt)
 
 DocMeta.setdocmeta!(GeometricFlux, :DocTestSetup, :(using GeometricFlux, Flux); recursive=true)
 
+# DemoCards
+demopage, postprocess_cb, demo_assets = makedemos("tutorials")
+isnothing(demo_assets) || (push!(ASSETS, demo_assets))
+
 makedocs(
     bib,
     sitename = "GeometricFlux.jl",
     format = Documenter.HTML(
-      assets = ["assets/flux.css", "assets/favicon.ico"],
+      assets = ASSETS,
       canonical = "https://fluxml.ai/GeometricFlux.jl/stable/",
       analytics = "G-M61P0B2Y8E",
+      edit_link = "master",
     ),
     clean = false,
     modules = [GeometricFlux,GraphSignals],
     pages = ["Home" => "index.md",
-             "Tutorials" => [
-                 "Semi-Supervised Learning with GCN" => "tutorials/semisupervised_gcn.md",
-                 "GCN with Fixed Graph" => "tutorials/gcn_fixed_graph.md",
-                 "Graph Attention Network" => "tutorials/gat.md",
-                 "DeepSet for Digit Sum" => "tutorials/deepset.md",
-                 "Variational Graph Autoencoder" => "tutorials/vgae.md",
-                 "Graph Embedding" => "tutorials/graph_embedding.md",
-              ],
+             demopage,
              "Introduction" => "introduction.md",
              "Basics" => [
                  "Graph Convolutions" => "basics/conv.md",
@@ -55,6 +56,9 @@ makedocs(
     ]
 )
 
+# callbacks of DemoCards
+postprocess_cb()
+
 deploydocs(
   repo = "github.com/FluxML/GeometricFlux.jl.git",
   target = "build",

data/cora_labels.jld2

@@ -0,0 +1,3 @@
+{
+  "theme": "grid"
+}

data/cora_paper2idx.jld2

@@ -0,0 +1,12 @@
+{
+  "theme": "grid",
+  "description": "To begin with GeometricFlux, it is recommended to learn with following examples.",
+  "order": [
+    "semisupervised_gcn.md",
+    "gcn_static_graph.md",
+    "gat.md",
+    "deepset.md",
+    "vgae.md",
+    "graph_embedding.md"
+  ]
+}

docs/Project.toml

@@ -1,4 +1,10 @@
-# Predicting Digits Sum from DeepSet model
+---
+title: Predicting Digits Sum from DeepSet Model
+cover: assets/logo.svg
+id: deepset
+---
+
+# Predicting Digits Sum from DeepSet Model
 
 Digits sum is a task of summing up digits in images or text. This example demonstrates summing up digits in arbitrary number of MNIST images. To accomplish such task, DeepSet model is suitable for this task. DeepSet model is excellent at the task which takes a set of objects and reduces them into single object.
 
@@ -9,8 +15,9 @@ Since a DeepSet model predicts the summation from a set of images, we have to pr
 First, the whole dataset is loaded from MLDatasets.jl and then shuffled before generating training dataset.
 
 ```julia
-train_X, train_y = MLDatasets.MNIST.traindata(Float32)
-train_X, train_y = shuffle_data(train_X, train_y)
+train_data, test_data = MNIST(:train), MNIST(:test)
+train_X, train_y = shuffle_data(train_data.features, train_data.targets)
+test_X, test_y = shuffle_data(test_data.features, test_data.targets)
 ```
 
 The `generate_featuredgraphs` here generates a set of pairs which contains a `FeaturedGraph` and a summed number for prediction target. In a `FeaturedGraph`, an arbitrary number of MNIST images are collected as node features and corresponding nodes are collected in a graph without edges.
@@ -68,9 +75,8 @@ for epoch = 1:args.epochs
     @info "Epoch $(epoch)"
 
     for batch in train_loader
-        train_loss, back = Flux.pullback(ps) do
-            model_loss(model, batch |> device)
-        end
+        batch = batch |> device
+        train_loss, back = Flux.pullback(() -> model_loss(model, batch), ps)
         test_loss = model_loss(model, test_loader, device)
         grad = back(1f0)
         Flux.Optimise.update!(opt, ps, grad)

docs/bibliography.bib

@@ -1,3 +1,9 @@
+---
+title: Graph Attention Network
+cover: assets/logo.svg
+id: gat
+---
+
 # Graph Attention Network
 
 Graph attention network (GAT) belongs to the message-passing network family, and it queries node feature over its neighbor features and generates result as layer output.
@@ -7,18 +13,26 @@ Graph attention network (GAT) belongs to the message-passing network family, and
 We load dataset from Planetoid dataset. Here cora dataset is used.
 
 ```julia
-train_X, train_y = map(x -> Matrix(x), alldata(Planetoid(), dataset, padding=true))
+data = dataset[1].node_data
+X, y = data.features, onehotbatch(data.targets, 1:7)
+train_idx, test_idx = data.train_mask, data.val_mask
 ```
 
 ## Step 2: Batch up Features and Labels
 
 Just batch up features as usual.
 
 ```julia
+s, t = dataset[1].edge_index
+g = Graphs.Graph(dataset[1].num_nodes)
+for (i, j) in zip(s, t)
+    Graphs.add_edge!(g, i, j)
+end
+
 add_all_self_loops!(g)
 fg = FeaturedGraph(g)
-train_data = (repeat(train_X, outer=(1,1,train_repeats)), repeat(train_y, outer=(1,1,train_repeats)))
-train_loader = DataLoader(train_data, batchsize=batch_size, shuffle=true)
+train_X, train_y = repeat(X, outer=(1,1,train_repeats)), repeat(y, outer=(1,1,train_repeats))
+train_loader = DataLoader((train_X, train_y), batchsize=batch_size, shuffle=true)
 ```
 
 Notably, self loop for all nodes are needed for GAT model.
@@ -66,9 +80,8 @@ for epoch = 1:args.epochs
     @info "Epoch $(epoch)"
 
     for (X, y) in train_loader
-        loss, back = Flux.pullback(ps) do
-            model_loss(model, X |> device, y |> device, train_idx |> device)
-        end
+        X, y, device_idx = X |> device, y |> device, train_idx |> device
+        loss, back = Flux.pullback(() -> model_loss(model, X, y, device_idx), ps)
         train_acc = accuracy(model, train_loader, device, train_idx)
         test_acc = accuracy(model, test_loader, device, test_idx)
         grad = back(1f0)

docs/make.jl

@@ -1,10 +1,16 @@
-# GCN with Fixed Graph
+---
+title: GCN with Static Graph
+cover: assets/logo.svg
+id: gcn_static_graph
+---
 
-In the tutorial for semi-supervised learning with GCN, variable graphs are provided to GNN from `FeaturedGraph`, which contains a graph and node features. Each `FeaturedGraph` object can contain different graph and different node features, and can be train on the same GNN model. However, variable graph doesn't have the proper form of graph structure with respect to GNN layers and this lead to inefficient training/inference process. Fixed graph strategy can be used to train a GNN model with the same graph structure in GeometricFlux.
+# GCN with Static Graph
 
-## Fixed Graph
+In the tutorial for semi-supervised learning with GCN, variable graphs are provided to GNN from `FeaturedGraph`, which contains a graph and node features. Each `FeaturedGraph` object can contain different graph and different node features, and can be train on the same GNN model. However, variable graph doesn't have the proper form of graph structure with respect to GNN layers and this lead to inefficient training/inference process. Static graph strategy can be used to train a GNN model with the same graph structure in GeometricFlux.
 
-A fixed graph is given to a layer by `WithGraph` syntax. `WithGraph` wrap a `FeaturedGraph` object and a GNN layer as first and second arguments, respectively.
+## Static Graph
+
+A static graph is given to a layer by `WithGraph` syntax. `WithGraph` wrap a `FeaturedGraph` object and a GNN layer as first and second arguments, respectively.
 
 ```julia
 fg = FeaturedGraph(graph)
@@ -26,23 +32,29 @@ Since features are in the form of array, they can be batched up for batched lear
 Different from loading datasets in semi-supervised learning example, we use `alldata` for supervised learning here and `padding=true` is added in order to padding features from partial nodes to pseudo-full nodes. A padded features contains zeros in the nodes that are not supposed to be train on.
 
 ```julia
-train_X, train_y = map(x -> Matrix(x), alldata(Planetoid(), dataset, padding=true))
+data = dataset[1].node_data
+X, y = data.features, onehotbatch(data.targets, 1:7)
+train_idx, test_idx = data.train_mask, data.val_mask
+train_X, train_y = repeat(X, outer=(1,1,train_repeats)), repeat(y, outer=(1,1,train_repeats))
 ```
 
 We need graph and node indices for training as well.
 
 ```julia
-g = graphdata(Planetoid(), dataset)
-train_idx = 1:size(train_X, 2)
+s, t = dataset[1].edge_index
+g = Graphs.Graph(dataset[1].num_nodes)
+for (i, j) in zip(s, t)
+    Graphs.add_edge!(g, i, j)
+end
+fg = FeaturedGraph(g)
 ```
 
 ## Step 2: Batch up Features and Labels
 
 In order to make batch learning available, we separate graph and node features. We don't subgraph here. Node features are batched up by repeating node features here for demonstration, since planetoid dataset doesn't have batched settings. Different repeat numbers can be specified by `train_repeats` and `train_repeats`.
 
 ```julia
-fg = FeaturedGraph(g)
-train_data = (repeat(train_X, outer=(1,1,train_repeats)), repeat(train_y, outer=(1,1,train_repeats)))
+train_loader = DataLoader((train_X, train_y), batchsize=batch_size, shuffle=true)
 ```
 
 ## Step 3: Build a GCN model
@@ -99,7 +111,8 @@ for epoch = 1:args.epochs
     @info "Epoch $(epoch)"
 
     for (X, y) in train_loader
-        grad = gradient(() -> model_loss(model, args.λ, X |> device, y |> device, train_idx |> device), ps)
+        X, y, device_idx = X |> device, y |> device, train_idx |> device
+        grad = gradient(() -> model_loss(model, args.λ, X, y, device_idx), ps)
         Flux.Optimise.update!(opt, ps, grad)
         train_steps += 1
     end

docs/src/tutorials/graph_embedding.md

@@ -0,0 +1,7 @@
+---
+title: Graph Embedding Through Node2vec Model
+cover: assets/logo.svg
+id: graph_embedding
+---
+
+# Graph Embedding Through Node2vec Model

docs/tutorials/config.json

@@ -1,3 +1,9 @@
+---
+title: Semi-supervised Learning with Graph Convolution Networks (GCN)
+cover: assets/logo.svg
+id: semisupervised_gcn
+---
+
 # Semi-supervised Learning with Graph Convolution Networks (GCN)
 
 Graph convolution networks (GCN) have been considered as the first step to graph neural networks (GNN). This example will go through how to train a vanilla GCN.

docs/tutorials/examples/assets/logo.svg

@@ -1,3 +1,9 @@
+---
+title: Variational Graph Autoencoder
+cover: assets/logo.svg
+id: vgae
+---
+
 # Variational Graph Autoencoder
 
 Variational Graph Autoencoder (VGAE) is a unsupervised generative model. It takes node features and graph structure and predicts the edge link in the graph. A link preidction task is defined for this model.
@@ -7,13 +13,19 @@ Variational Graph Autoencoder (VGAE) is a unsupervised generative model. It take
 We load dataset from Planetoid dataset. Here cora dataset is used.
 
 ```julia
-train_X, _ = map(x -> Matrix(x), alldata(Planetoid(), dataset))
+data = dataset[1].node_data
+X = data.features
+train_X = repeat(X, outer=(1, 1, train_repeats))
 ```
 
 Notably, a link prediction task will output a graph in the form of adjacency matrix, so an adjacency matrix is needed as label for this task.
 
 ```julia
-g = graphdata(Planetoid(), dataset)
+s, t = dataset[1].edge_index
+g = Graphs.Graph(dataset[1].num_nodes)
+for (i, j) in zip(s, t)
+    Graphs.add_edge!(g, i, j)
+end
 fg = FeaturedGraph(g)
 A = GraphSignals.adjacency_matrix(fg)
 ```
@@ -23,8 +35,7 @@ A = GraphSignals.adjacency_matrix(fg)
 Just batch up features as usual.
 
 ```julia
-data = (repeat(X, outer=(1,1,train_repeats)), repeat(A, outer=(1,1,train_repeats)))
-loader = DataLoader(data, batchsize=batch_size, shuffle=true)
+loader = DataLoader((train_X, train_y), batchsize=batch_size, shuffle=true)
 ```
 
 ## Step 3: Build a VGAE model
@@ -90,10 +101,9 @@ ps = Flux.params(model)
 for epoch = 1:args.epochs
     @info "Epoch $(epoch)"
 
-    for (X, A) in loader
-        loss, back = Flux.pullback(ps) do
-            model_loss(model, X |> device, A |> device, args.β)
-        end
+    for (X, Â) in loader
+        X, Â = X |> device, Â |> device
+        loss, back = Flux.pullback(() -> model_loss(model, X, Â, args.β), ps)
         prec = precision(model, loader, device)
         grad = back(1f0)
         Flux.Optimise.update!(opt, ps, grad)

docs/tutorials/examples/config.json

@@ -0,0 +1,3 @@
+# [Tutorials](@id tutorials)
+
+{{{democards}}}

docs/src/tutorials/deepset.md renamed to docs/tutorials/examples/deepset.md

@@ -20,19 +20,15 @@ function load_data(
     test_min_length,
     test_max_length
 )
-    train_X, train_y = MLDatasets.MNIST.traindata(Float32)
-    test_X, test_y = MLDatasets.MNIST.testdata(Float32)
-
-    train_X, train_y = shuffle_data(train_X, train_y)
-    test_X, test_y = shuffle_data(test_X, test_y)
+    train_data, test_data = MNIST(:train), MNIST(:test)
+    train_X, train_y = shuffle_data(train_data.features, train_data.targets)
+    test_X, test_y = shuffle_data(test_data.features, test_data.targets)
 
     train_data = generate_featuredgraphs(train_X, train_y, num_train_examples, 1:train_max_length)
     test_data = generate_featuredgraphs(test_X, test_y, num_test_examples, test_min_length:test_max_length)
-    train_batch = Flux.batch(train_data)
-    test_batch = Flux.batch(test_data)
 
-    train_loader = DataLoader(train_batch, batchsize=batch_size)
-    test_loader = DataLoader(test_batch, batchsize=batch_size)
+    train_loader = DataLoader(train_data, batchsize=batch_size)
+    test_loader = DataLoader(test_data, batchsize=batch_size)
     return train_loader, test_loader
 end
 
@@ -92,6 +88,7 @@ function train(; kws...)
     # GPU config
     if args.cuda && CUDA.has_cuda()
         device = gpu
+        CUDA.allowscalar(false)
         @info "Training on GPU"
     else
         device = cpu
@@ -131,9 +128,8 @@ function train(; kws...)
         progress = Progress(length(train_loader))
 
         for batch in train_loader
-            train_loss, back = Flux.pullback(ps) do
-                model_loss(model, batch |> device)
-            end
+            batch = batch |> device
+            train_loss, back = Flux.pullback(() -> model_loss(model, batch), ps)
             test_loss = model_loss(model, test_loader, device)
             grad = back(1f0)
             Flux.Optimise.update!(opt, ps, grad)