update examples to explicit gradient (#301)

* update to explicit gradient

* update DiffEqFlux

* regenerate pluto output

* change test

Author: CarloLucibello

.github/workflows/docs.yml

@@ -14,7 +14,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@latest
         with:
-          version: '1.8.2'
+          version: '1.9.1'
       - name: Install dependencies
         run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy

README.md

@@ -19,6 +19,7 @@ Among its features:
 * CUDA support.
 * Integration with [Graphs.jl](https://github.com/JuliaGraphs/Graphs.jl).
 * [Examples](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/tree/master/examples) of node, edge, and graph level machine learning tasks. 
+* Heterogeneous and temporal graphs. 
 
 ## Installation
 
@@ -30,7 +31,7 @@ pkg> add GraphNeuralNetworks
 
 ## Usage
 
-Usage examples can be found in the [examples](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/tree/master/examples) and in the [notebooks](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/tree/master/notebooks)  folder. Also, make sure to read the [documentation](https://CarloLucibello.github.io/GraphNeuralNetworks.jl/dev) for a comprehensive introduction to the library.
+Usage examples can be found in the [examples](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/tree/master/examples) and in the [notebooks](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/tree/master/notebooks) folder. Also, make sure to read the [documentation](https://CarloLucibello.github.io/GraphNeuralNetworks.jl/dev) for a comprehensive introduction to the library.
 
 
 ## Citing

docs/Project.toml

@@ -7,10 +7,11 @@ Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MarkdownLiteral = "736d6165-7244-6769-4267-6b50796e6954"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781"
+PlutoStaticHTML = "359b1769-a58e-495b-9770-312e911026ad"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
-
 [compat]
-DemoCards = "^0.4.11"
+DemoCards = "0.5.0"

docs/make.jl

@@ -1,4 +1,5 @@
 using Flux, NNlib, GraphNeuralNetworks, Graphs, SparseArrays
+using Pluto, PlutoStaticHTML # for tutorials
 using Documenter, DemoCards
 
 tutorials, tutorials_cb, tutorial_assets = makedemos("tutorials")

docs/pluto_output/gnn_intro_pluto.md

@@ -25,8 +25,8 @@
 <!--
     # This information is used for caching.
     [PlutoStaticHTML.State]
-    input_sha = "465a4590b716e99f2d70eaaa9ca2e5de5f8549d401b4fa432fab439ab30aa172"
-    julia_version = "1.8.2"
+    input_sha = "f145b80b8f1e399d4cd5686b529cf173942102c538702952fe0743defca62210"
+    julia_version = "1.9.1"
 -->
 <pre class='language-julia'><code class='language-julia'>begin
     using Flux
@@ -102,18 +102,18 @@ end</code></pre>
 <div class="markdown"><p>We have some useful utilities for working with graph datasets, <em>e.g.</em>, we can shuffle the dataset and use the first 150 graphs as training graphs, while using the remaining ones for testing:</p></div>
 
 <pre class='language-julia'><code class='language-julia'>train_data, test_data = splitobs((graphs, y), at = 150, shuffle = true) |&gt; getobs</code></pre>
-<pre class="code-output documenter-example-output" id="var-train_data">((GraphNeuralNetworks.GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}[GNNGraph(16, 34), GNNGraph(15, 32), GNNGraph(19, 44), GNNGraph(20, 44), GNNGraph(20, 46), GNNGraph(15, 34), GNNGraph(18, 40), GNNGraph(16, 36), GNNGraph(13, 28), GNNGraph(16, 34)  …  GNNGraph(23, 48), GNNGraph(20, 44), GNNGraph(28, 66), GNNGraph(25, 56), GNNGraph(13, 28), GNNGraph(16, 36), GNNGraph(12, 24), GNNGraph(22, 50), GNNGraph(25, 58), GNNGraph(19, 42)], Bool[1 1 … 0 0; 0 0 … 1 1]), (GraphNeuralNetworks.GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}[GNNGraph(12, 24), GNNGraph(11, 22), GNNGraph(15, 34), GNNGraph(19, 44), GNNGraph(22, 50), GNNGraph(17, 38), GNNGraph(17, 38), GNNGraph(17, 38), GNNGraph(19, 42), GNNGraph(13, 28)  …  GNNGraph(19, 40), GNNGraph(13, 28), GNNGraph(22, 50), GNNGraph(14, 28), GNNGraph(23, 54), GNNGraph(20, 46), GNNGraph(13, 28), GNNGraph(26, 60), GNNGraph(17, 38), GNNGraph(12, 26)], Bool[1 1 … 0 0; 0 0 … 1 1]))</pre>
+<pre class="code-output documenter-example-output" id="var-train_data">((GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}[GNNGraph(12, 24) with x: 7×12 data, GNNGraph(22, 50) with x: 7×22 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(25, 56) with x: 7×25 data, GNNGraph(16, 36) with x: 7×16 data, GNNGraph(11, 22) with x: 7×11 data, GNNGraph(18, 38) with x: 7×18 data, GNNGraph(23, 52) with x: 7×23 data, GNNGraph(22, 50) with x: 7×22 data, GNNGraph(20, 46) with x: 7×20 data  …  GNNGraph(16, 34) with x: 7×16 data, GNNGraph(13, 28) with x: 7×13 data, GNNGraph(21, 44) with x: 7×21 data, GNNGraph(17, 38) with x: 7×17 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(12, 24) with x: 7×12 data, GNNGraph(22, 50) with x: 7×22 data, GNNGraph(19, 42) with x: 7×19 data, GNNGraph(16, 34) with x: 7×16 data, GNNGraph(16, 36) with x: 7×16 data], Bool[1 0 … 1 0; 0 1 … 0 1]), (GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}[GNNGraph(21, 44) with x: 7×21 data, GNNGraph(22, 50) with x: 7×22 data, GNNGraph(16, 34) with x: 7×16 data, GNNGraph(27, 66) with x: 7×27 data, GNNGraph(13, 26) with x: 7×13 data, GNNGraph(20, 44) with x: 7×20 data, GNNGraph(19, 44) with x: 7×19 data, GNNGraph(20, 46) with x: 7×20 data, GNNGraph(16, 34) with x: 7×16 data, GNNGraph(13, 28) with x: 7×13 data  …  GNNGraph(11, 22) with x: 7×11 data, GNNGraph(20, 46) with x: 7×20 data, GNNGraph(16, 34) with x: 7×16 data, GNNGraph(18, 40) with x: 7×18 data, GNNGraph(13, 28) with x: 7×13 data, GNNGraph(20, 44) with x: 7×20 data, GNNGraph(14, 30) with x: 7×14 data, GNNGraph(13, 26) with x: 7×13 data, GNNGraph(21, 44) with x: 7×21 data, GNNGraph(22, 50) with x: 7×22 data], Bool[0 0 … 0 0; 1 1 … 1 1]))</pre>
 
 <pre class='language-julia'><code class='language-julia'>begin
-    train_loader = DataLoader(train_data, batchsize = 64, shuffle = true)
-    test_loader = DataLoader(test_data, batchsize = 64, shuffle = false)
+    train_loader = DataLoader(train_data, batchsize = 32, shuffle = true)
+    test_loader = DataLoader(test_data, batchsize = 32, shuffle = false)
 end</code></pre>
-<pre class="code-output documenter-example-output" id="var-test_loader">1-element DataLoader(::Tuple{Vector{GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}}, OneHotArrays.OneHotMatrix{UInt32, Vector{UInt32}}}, batchsize=64)
+<pre class="code-output documenter-example-output" id="var-test_loader">2-element DataLoader(::Tuple{Vector{GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}}, OneHotArrays.OneHotMatrix{UInt32, Vector{UInt32}}}, batchsize=32)
   with first element:
-  (38-element Vector{GraphNeuralNetworks.GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}}, 2×38 OneHotMatrix(::Vector{UInt32}) with eltype Bool,)</pre>
+  (32-element Vector{GraphNeuralNetworks.GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}}, 2×32 OneHotMatrix(::Vector{UInt32}) with eltype Bool,)</pre>
 
 
-<div class="markdown"><p>Here, we opt for a <code>batch_size</code> of 64, leading to 3 (randomly shuffled) mini-batches, containing all <span class="tex">$2 \cdot 64+22 = 150$</span> graphs.</p></div>
+<div class="markdown"><p>Here, we opt for a <code>batch_size</code> of 32, leading to 5 (randomly shuffled) mini-batches, containing all <span class="tex">$4 \cdot 32+22 = 150$</span> graphs.</p></div>
 
 
 ```
@@ -123,15 +123,15 @@ end</code></pre>
 <p>Since graphs in graph classification datasets are usually small, a good idea is to <strong>batch the graphs</strong> before inputting them into a Graph Neural Network to guarantee full GPU utilization. In the image or language domain, this procedure is typically achieved by <strong>rescaling</strong> or <strong>padding</strong> each example into a set of equally-sized shapes, and examples are then grouped in an additional dimension. The length of this dimension is then equal to the number of examples grouped in a mini-batch and is typically referred to as the <code>batchsize</code>.</p><p>However, for GNNs the two approaches described above are either not feasible or may result in a lot of unnecessary memory consumption. Therefore, GraphNeuralNetworks.jl opts for another approach to achieve parallelization across a number of examples. Here, adjacency matrices are stacked in a diagonal fashion (creating a giant graph that holds multiple isolated subgraphs), and node and target features are simply concatenated in the node dimension (the last dimension).</p><p>This procedure has some crucial advantages over other batching procedures:</p><ol><li><p>GNN operators that rely on a message passing scheme do not need to be modified since messages are not exchanged between two nodes that belong to different graphs.</p></li><li><p>There is no computational or memory overhead since adjacency matrices are saved in a sparse fashion holding only non-zero entries, <em>i.e.</em>, the edges.</p></li></ol><p>GraphNeuralNetworks.jl can <strong>batch multiple graphs into a single giant graph</strong>:</p></div>
 
 <pre class='language-julia'><code class='language-julia'>vec_gs, _ = first(train_loader)</code></pre>
-<pre class="code-output documenter-example-output" id="var-vec_gs">(GraphNeuralNetworks.GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}[GNNGraph(11, 22), GNNGraph(16, 36), GNNGraph(16, 34), GNNGraph(22, 50), GNNGraph(18, 40), GNNGraph(19, 40), GNNGraph(24, 50), GNNGraph(23, 54), GNNGraph(24, 50), GNNGraph(12, 26)  …  GNNGraph(20, 44), GNNGraph(11, 22), GNNGraph(22, 50), GNNGraph(13, 26), GNNGraph(16, 34), GNNGraph(10, 20), GNNGraph(28, 66), GNNGraph(19, 44), GNNGraph(14, 30), GNNGraph(18, 38)], Bool[1 0 … 1 0; 0 1 … 0 1])</pre>
+<pre class="code-output documenter-example-output" id="var-vec_gs">(GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}[GNNGraph(17, 38) with x: 7×17 data, GNNGraph(19, 42) with x: 7×19 data, GNNGraph(13, 28) with x: 7×13 data, GNNGraph(14, 30) with x: 7×14 data, GNNGraph(13, 28) with x: 7×13 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(16, 36) with x: 7×16 data, GNNGraph(24, 50) with x: 7×24 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(15, 34) with x: 7×15 data  …  GNNGraph(16, 34) with x: 7×16 data, GNNGraph(16, 34) with x: 7×16 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(12, 26) with x: 7×12 data, GNNGraph(17, 38) with x: 7×17 data, GNNGraph(20, 44) with x: 7×20 data, GNNGraph(13, 28) with x: 7×13 data, GNNGraph(26, 60) with x: 7×26 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(24, 50) with x: 7×24 data], Bool[0 0 … 0 0; 1 1 … 1 1])</pre>
 
 <pre class='language-julia'><code class='language-julia'>MLUtils.batch(vec_gs)</code></pre>
 <pre class="code-output documenter-example-output" id="var-hash102363">GNNGraph:
-    num_nodes = 1191
-    num_edges = 2618
-    num_graphs = 64
-    ndata:
-        x =&gt; 7×1191 Matrix{Float32}</pre>
+  num_nodes: 585
+  num_edges: 1292
+  num_graphs: 32
+  ndata:
+	x = 7×585 Matrix{Float32}</pre>
 
 
 <div class="markdown"><p>Each batched graph object is equipped with a <strong><code>graph_indicator</code> vector</strong>, which maps each node to its respective graph in the batch:</p><p class="tex">$$\textrm{graph\_indicator} = [1, \ldots, 1, 2, \ldots, 2, 3, \ldots ]$$</p></div>
@@ -177,8 +177,7 @@ end</code></pre>
     # device = Flux.gpu # uncomment this for GPU training
     device = Flux.cpu
     model = model |&gt; device
-    ps = Flux.params(model)
-    opt = Adam(1e-3)
+    opt = Flux.setup(Adam(1e-3), model)
 
     function report(epoch)
         train = eval_loss_accuracy(model, train_loader, device)
@@ -190,11 +189,11 @@ end</code></pre>
     for epoch in 1:epochs
         for (g, y) in train_loader
             g, y = MLUtils.batch(g) |&gt; device, y |&gt; device
-            gs = Flux.gradient(ps) do
+            grad = Flux.gradient(model) do model
                 ŷ = model(g, g.ndata.x)
                 logitcrossentropy(ŷ, y)
             end
-            Flux.Optimise.update!(opt, ps, gs)
+            Flux.update!(opt, model, grad[1])
         end
         epoch % infotime == 0 && report(epoch)
     end