ok works

Author: aurorarossi

GNNLux/docs/src/index.md

@@ -12,7 +12,7 @@ We create a dataset consisting in multiple random graphs and associated data fea
 
 ```julia
 using GNNLux, Lux, Statistics, MLUtils, Random
-using Zygote, Optimizers
+using Zygote, Optimisers
 
 all_graphs = GNNGraph[]
 
@@ -22,61 +22,49 @@ for _ in 1:1000
             gdata=(; y = randn(Float32)))         # Regression target   
     push!(all_graphs, g)
 end
-```
 
-### Model building 
+train_graphs, test_graphs = MLUtils.splitobs(all_graphs, at=0.8)
 
-We concisely define our model as a [`GNNLux.GNNChain`](@ref) containing two graph convolutional layers. If CUDA is available, our model will live on the gpu.
+# g = rand_graph(10, 40, ndata=(; x = randn(Float32, 16,10)), gdata=(; y = randn(Float32))) 
 
-```julia
-device = CUDA.functional() ? Lux.gpu_device() : Lux.cpu_device()
 rng = Random.default_rng()
 
 model = GNNChain(GCNConv(16 => 64),
                 x -> relu.(x),     
                 GCNConv(64 => 64, relu),
-                GlobalMeanPool(),  # Aggregate node-wise features into graph-wise features
+                x -> mean(x, dims=2),
                 Dense(64, 1)) 
 
 ps, st = LuxCore.setup(rng, model)
-```
-
-### Training 
 
-
-```julia
-train_graphs, test_graphs = MLUtils.splitobs(all_graphs, at=0.8)
-
-train_loader = MLUtils.DataLoader(train_graphs, 
-                batchsize=32, shuffle=true, collate=true)
-test_loader = MLUtils.DataLoader(test_graphs, 
-                batchsize=32, shuffle=false, collate=true)
-
-for epoch in 1:100
-    for g in train_loader
-        g = g |> device
-        grad = gradient(model -> loss(model, g), model)
-        Flux.update!(opt, model, grad[1])
-    end
-
-    @info (; epoch, train_loss=loss(model, train_loader), test_loss=loss(model, test_loader))
+function custom_loss(model, ps,st,tuple)
+    g,x,y = tuple
+    y_pred,st = model(g, x, ps, st)  
+    return MSELoss()(y_pred, y), (layers = st,), 0
 end
 
-function train_model!(model, ps, st, train_loader)
-    train_state = Lux.Training.TrainState(model, ps, st, Adam(0.001f0))
 
-    for iter in 1:1000
-        for g in train_loader
-            _, loss, _, train_state = Lux.Training.single_train_step!(AutoZygote(), MSELoss(),
-                ((g, g.x)...,g.y), train_state)
-            if iter % 100 == 1 || iter == 1000
-                @info "Iteration: %04d \t Loss: %10.9g\n" iter loss
+function train_model!(model, ps, st, train_graphs, test_graphs)
+    train_state = Lux.Training.TrainState(model, ps, st, Adam(0.0001f0))
+    loss=0
+    for iter in 1:100
+        for g in train_graphs
+            _, loss, _, train_state = Lux.Training.single_train_step!(AutoZygote(), custom_loss,(g, g.x, g.y), train_state)
+        end
+        if iter % 10 == 0 || iter == 100
+            st_ = Lux.testmode(train_state.states)
+            test_loss =0
+            for g in test_graphs
+                ŷ, st_ = model(g, g.x, train_state.parameters, st_)
+                st_ = (layers = st_,)
+                test_loss += MSELoss()(g.y,ŷ)
             end
+            test_loss = test_loss/length(test_graphs)
+            @info (; iter, loss, test_loss)
         end
     end
 
     return model, ps, st
 end
 
-train_model!(model, ps, st, train_loader)
-```
+train_model!(model, ps, st, train_graphs, test_graphs)