EvolveGCNOCell

CarloLucibello · CarloLucibello · commit 383d46444b11 · 2024-12-16T12:19:09.000+01:00
diff --git a/GraphNeuralNetworks/src/GraphNeuralNetworks.jl b/GraphNeuralNetworks/src/GraphNeuralNetworks.jl
@@ -54,9 +54,9 @@ export GNNRecurrence,
        GConvGRU, GConvGRUCell,
        GConvLSTM, GConvLSTMCell,
        DCGRU, DCGRUCell,
+       EvolveGCNO, EvolveGCNOCell,
        TGCN,
-       A3TGCN,
-       EvolveGCNO
+       A3TGCN
 
 include("layers/pool.jl")
 export GlobalPool,
diff --git a/GraphNeuralNetworks/src/layers/temporalconv.jl b/GraphNeuralNetworks/src/layers/temporalconv.jl
@@ -1,12 +1,23 @@
 function scan(cell, g::GNNGraph, x::AbstractArray{T,3}, state) where {T}
     y = []
-    for x_t in eachslice(x, dims = 2)
-        yt, state = cell(g, x_t, state)
+    for xt in eachslice(x, dims = 2)
+        yt, state = cell(g, xt, state)
         y = vcat(y, [yt])
     end
     return stack(y, dims = 2)
 end
 
+function scan(cell, tg::TemporalSnapshotsGNNGraph, x::AbstractVector, state)
+    # @assert length(x) == length(tg.snapshots)
+    y = []
+    for (t, xt) in enumerate(x)
+        gt = tg.snapshots[t]
+        yt, state = cell(gt, xt, state)
+        y = vcat(y, [yt])
+    end
+    return y
+end
+
 
 """
     GNNRecurrence(cell)
@@ -20,7 +31,7 @@ to process an entire temporal sequence of node features at once.
 
 - `g`: The input graph.
 - `x`: The time-varying node features. An array of size `in x timesteps x num_nodes`.
-- `state`: The initial state of the cell. 
+- `state`: The current state of the cell. 
    If not provided, it is generated by calling `Flux.initialstates(cell)`.
 
 Applies the recurrent cell to each timestep of the input sequence and returns the output as
@@ -61,11 +72,11 @@ Flux.@layer GNNRecurrence
 
 Flux.initialstates(rnn::GNNRecurrence) = Flux.initialstates(rnn.cell)
 
-function (rnn::GNNRecurrence)(g::GNNGraph, x::AbstractArray{T,3}) where {T}
+function (rnn::GNNRecurrence)(g, x)
     return rnn(g, x, initialstates(rnn))
 end
 
-function (rnn::GNNRecurrence)(g::GNNGraph, x::AbstractArray{T,3}, state) where {T}
+function (rnn::GNNRecurrence)(g, x, state) where {T}
     return scan(rnn.cell, g, x, state)
 end
 
@@ -97,7 +108,7 @@ followed by a Gated Recurrent Unit (GRU) cell to model temporal dependencies.
 
 - `g`: The input graph.
 - `x`: The node features. It should be a matrix of size `in x num_nodes`.
-- `h`: The initial hidden state of the GRU cell. If given, it is a matrix of size `out x num_nodes`.
+- `h`: The current hidden state of the GRU cell. If given, it is a matrix of size `out x num_nodes`.
        If not provided, it is assumed to be a matrix of zeros.
 
 Performs one recurrence step and returns a tuple `(h, h)`, 
@@ -251,9 +262,9 @@ followed by a Long Short-Term Memory (LSTM) cell to model temporal dependencies.
 
 - `g`: The input graph.
 - `x`: The node features. It should be a matrix of size `in x num_nodes`.
-- `state`: The initial hidden state of the LSTM cell.  
+- `state`: The current state of the LSTM cell.  
        If given, it is a tuple `(h, c)` where both `h` and `c` are arrays of size `out x num_nodes`.
-       If not provided, the initial hidden state is assumed to be a tuple of matrices of zeros.
+       If not provided, it is assumed to be a tuple of matrices of zeros.
 
 Performs one recurrence step and returns a tuple `(output, state)`, 
 where `output` is the updated hidden state `h` of the LSTM cell and `state` is the updated tuple `(h, c)`.
@@ -434,7 +445,7 @@ in combination with a Gated Recurrent Unit (GRU) cell to model temporal dependen
 
 - `g`: The input graph.
 - `x`: The node features. It should be a matrix of size `in x num_nodes`.
-- `h`: The initial hidden state of the GRU cell. If given, it is a matrix of size `out x num_nodes`.
+- `h`: The current state of the GRU cell. It is a matrix of size `out x num_nodes`.
        If not provided, it is assumed to be a matrix of zeros.
 
 Performs one recurrence step and returns a tuple `(h, h)`,
@@ -547,4 +558,96 @@ julia> size(y) # (d_out, timesteps, num_nodes)
 """
 DCGRU(args...; kws...) = GNNRecurrence(DCGRUCell(args...; kws...))
 
+""""
+    EvolveGCNOCell(in => out; bias = true, init = glorot_uniform)
+
+Evolving Graph Convolutional Network cell of type "-O" from the paper 
+[EvolveGCN: Evolving Graph Convolutional Networks for Dynamic Graphs](https://arxiv.org/abs/1902.10191).
+
+Uses a [`GCNConv`](@ref) layer to model spatial dependencies, and an `LSTMCell` to model temporal dependencies.
+Can work with time-varying graphs and node features.
+
+# Arguments
+
+- `in => out`: A pair where `in` is the number of input node features and `out` 
+  is the number of output node features.
+- `bias`: Add learnable bias for the convolution and the lstm cell. Default `true`.
+- `init`: Weights' initializer for the convolution. Default `glorot_uniform`.
+
+# Forward 
+
+    cell(g::GNNGraph, x, [state]) -> x, state
+
+- `g`: The input graph.
+- `x`: The node features. It should be a matrix of size `in x num_nodes`.
+- `state`: The current state of the cell.
+    A state is a tuple `(weight, lstm)` where `weight` is the convolution's weight and `lstm` is the lstm's state.
+    If not provided, it is generated by calling `Flux.initialstates(cell)`.
+
+Returns the updated node features `x` and the updated state.
+
+```jldoctest
+julia> using GraphNeuralNetworks, Flux
+
+julia> num_nodes, num_edges = 5, 10;
+
+julia> d_in, d_out = 2, 3;
+
+julia> timesteps = 5;
+
+julia> g = [rand_graph(num_nodes, num_edges) for t in 1:timesteps];
+
+julia> x = [rand(Float32, d_in, num_nodes) for t in 1:timesteps];
+
+julia> cell1 = EvolveGCNOCell(d_in => d_out)
+EvolveGCNOCell(2 => 3)  # 321 parameters
+
+julia> cell2 = EvolveGCNOCell(d_out => d_out)
+EvolveGCNOCell(3 => 3)  # 696 parameters
+
+julia> state1 = Flux.initialstates(cell1);
+
+julia> state2 = Flux.initialstates(cell2);
+
+julia> outputs = [];
+
+julia> for t in 1:timesteps
+           zt, state1 = cell1(g[t], x[t], state1)
+           yt, state2 = cell2(g[t], zt, state2)
+           outputs = vcat(outputs, [yt])
+       end
+
+julia> size(outputs[end]) # (d_out, num_nodes)
+(3, 5)
+```
+"""
+struct EvolveGCNOCell{C,L} <: GNNLayer
+    in::Int
+    out::Int
+    conv::C
+    lstm::L
+end
+
+Flux.@layer :noexpand EvolveGCNOCell
 
+function EvolveGCNOCell((in,out)::Pair{Int,Int}; bias = true, init = glorot_uniform)
+    conv = GCNConv(in => out; bias, init)
+    lstm = LSTMCell(in*out => in*out; bias)
+    return EvolveGCNOCell(in, out, conv, lstm)
+end
+
+function Flux.initialstates(cell::EvolveGCNOCell)
+    weight = reshape(cell.conv.weight, :)
+    lstm = Flux.initialstates(cell.lstm)
+    return (; weight, lstm)
+end
+
+function (cell::EvolveGCNOCell)(g::GNNGraph, x::AbstractMatrix, state)
+    weight, state_lstm = cell.lstm(state.weight, state.lstm)
+    x = cell.conv(g, x, conv_weight = reshape(weight, (cell.out, cell.in)))
+    return x, (; weight, lstm = state_lstm)
+end
+
+function Base.show(io::IO, egcno::EvolveGCNOCell)
+    print(io, "EvolveGCNOCell($(egcno.in) => $(egcno.out))")
+end
diff --git a/GraphNeuralNetworks/src/layers/temporalconv_old.jl b/GraphNeuralNetworks/src/layers/temporalconv_old.jl
@@ -1,13 +1,3 @@
-function scan(cell, g, x, state)
-    y = []
-    for x_t in eachslice(x, dims = 2)
-        yt, state = cell(g, x_t, state)
-        y = vcat(y, [yt])
-    end
-    return stack(y, dims = 2)
-end
-
-
 struct TGCNCell{C,G} <: GNNLayer
     conv::C
     gru::G
@@ -179,97 +169,97 @@ function Base.show(io::IO, a3tgcn::A3TGCN)
 end
 
 
-# """
-#     EvolveGCNO(ch; bias = true, init = glorot_uniform, init_state = Flux.zeros32)
-
-# Evolving Graph Convolutional Network (EvolveGCNO) layer from the paper [EvolveGCN: Evolving Graph Convolutional Networks for Dynamic Graphs](https://arxiv.org/pdf/1902.10191).
-
-# Perfoms a Graph Convolutional layer with parameters derived from a Long Short-Term Memory (LSTM) layer across the snapshots of the temporal graph.
-
-
-# # Arguments
-
-# - `in`: Number of input features.
-# - `out`: Number of output features.
-# - `bias`: Add learnable bias. Default `true`.
-# - `init`: Weights' initializer. Default `glorot_uniform`.
-# - `init_state`: Initial state of the hidden stat of the LSTM layer. Default `zeros32`.
-
-# # Examples
-
-# ```jldoctest
-# julia> tg = TemporalSnapshotsGNNGraph([rand_graph(10,20; ndata = rand(4,10)), rand_graph(10,14; ndata = rand(4,10)), rand_graph(10,22; ndata = rand(4,10))])
-# TemporalSnapshotsGNNGraph:
-#   num_nodes: [10, 10, 10]
-#   num_edges: [20, 14, 22]
-#   num_snapshots: 3
-
-# julia> ev = EvolveGCNO(4 => 5)
-# EvolveGCNO(4 => 5)
-
-# julia> size(ev(tg, tg.ndata.x))
-# (3,)
-
-# julia> size(ev(tg, tg.ndata.x)[1])
-# (5, 10)
-# ```
-# """
-# struct EvolveGCNO
-#     conv
-#     W_init
-#     init_state
-#     in::Int
-#     out::Int
-#     Wf
-#     Uf
-#     Bf
-#     Wi
-#     Ui
-#     Bi
-#     Wo
-#     Uo
-#     Bo
-#     Wc
-#     Uc
-#     Bc
-# end
+"""
+    EvolveGCNO(ch; bias = true, init = glorot_uniform, init_state = Flux.zeros32)
+
+Evolving Graph Convolutional Network (EvolveGCNO) layer from the paper [EvolveGCN: Evolving Graph Convolutional Networks for Dynamic Graphs](https://arxiv.org/pdf/1902.10191).
+
+Perfoms a Graph Convolutional layer with parameters derived from a Long Short-Term Memory (LSTM) layer across the snapshots of the temporal graph.
+
+
+# Arguments
+
+- `in`: Number of input features.
+- `out`: Number of output features.
+- `bias`: Add learnable bias. Default `true`.
+- `init`: Weights' initializer. Default `glorot_uniform`.
+- `init_state`: Initial state of the hidden stat of the LSTM layer. Default `zeros32`.
+
+# Examples
+
+```jldoctest
+julia> tg = TemporalSnapshotsGNNGraph([rand_graph(10,20; ndata = rand(4,10)), rand_graph(10,14; ndata = rand(4,10)), rand_graph(10,22; ndata = rand(4,10))])
+TemporalSnapshotsGNNGraph:
+  num_nodes: [10, 10, 10]
+  num_edges: [20, 14, 22]
+  num_snapshots: 3
+
+julia> ev = EvolveGCNO(4 => 5)
+EvolveGCNO(4 => 5)
+
+julia> size(ev(tg, tg.ndata.x))
+(3,)
+
+julia> size(ev(tg, tg.ndata.x)[1])
+(5, 10)
+```
+"""
+struct EvolveGCNO
+    conv
+    W_init
+    init_state
+    in::Int
+    out::Int
+    Wf
+    Uf
+    Bf
+    Wi
+    Ui
+    Bi
+    Wo
+    Uo
+    Bo
+    Wc
+    Uc
+    Bc
+end
  
-# function EvolveGCNO(ch; bias = true, init = glorot_uniform, init_state = Flux.zeros32)
-#     in, out = ch
-#     W = init(out, in)
-#     conv = GCNConv(ch; bias = bias, init = init)
-#     Wf = init(out, in)
-#     Uf = init(out, in)
-#     Bf = bias ? init(out, in) : nothing
-#     Wi = init(out, in)
-#     Ui = init(out, in)
-#     Bi = bias ? init(out, in) : nothing
-#     Wo = init(out, in)
-#     Uo = init(out, in)
-#     Bo = bias ? init(out, in) : nothing
-#     Wc = init(out, in)
-#     Uc = init(out, in)
-#     Bc = bias ? init(out, in) : nothing
-#     return EvolveGCNO(conv, W, init_state, in, out, Wf, Uf, Bf, Wi, Ui, Bi, Wo, Uo, Bo, Wc, Uc, Bc)
-# end
-
-# function (egcno::EvolveGCNO)(tg::TemporalSnapshotsGNNGraph, x)
-#     H = egcno.init_state(egcno.out, egcno.in)
-#     C = egcno.init_state(egcno.out, egcno.in)
-#     W = egcno.W_init
-#     X = map(1:tg.num_snapshots) do i
-#         F = Flux.sigmoid_fast.(egcno.Wf .* W + egcno.Uf .* H + egcno.Bf)
-#         I = Flux.sigmoid_fast.(egcno.Wi .* W + egcno.Ui .* H + egcno.Bi)
-#         O = Flux.sigmoid_fast.(egcno.Wo .* W + egcno.Uo .* H + egcno.Bo)
-#         C̃ = Flux.tanh_fast.(egcno.Wc .* W + egcno.Uc .* H + egcno.Bc)
-#         C = F .* C + I .* C̃
-#         H = O .* tanh_fast.(C)
-#         W = H
-#         egcno.conv(tg.snapshots[i], x[i]; conv_weight = H)
-#     end
-#     return X
-# end
+function EvolveGCNO(ch; bias = true, init = glorot_uniform, init_state = Flux.zeros32)
+    in, out = ch
+    W = init(out, in)
+    conv = GCNConv(ch; bias = bias, init = init)
+    Wf = init(out, in)
+    Uf = init(out, in)
+    Bf = bias ? init(out, in) : nothing
+    Wi = init(out, in)
+    Ui = init(out, in)
+    Bi = bias ? init(out, in) : nothing
+    Wo = init(out, in)
+    Uo = init(out, in)
+    Bo = bias ? init(out, in) : nothing
+    Wc = init(out, in)
+    Uc = init(out, in)
+    Bc = bias ? init(out, in) : nothing
+    return EvolveGCNO(conv, W, init_state, in, out, Wf, Uf, Bf, Wi, Ui, Bi, Wo, Uo, Bo, Wc, Uc, Bc)
+end
+
+function (egcno::EvolveGCNO)(tg::TemporalSnapshotsGNNGraph, x)
+    H = egcno.init_state(egcno.out, egcno.in)
+    C = egcno.init_state(egcno.out, egcno.in)
+    W = egcno.W_init
+    X = map(1:tg.num_snapshots) do i
+        F = Flux.sigmoid_fast.(egcno.Wf .* W + egcno.Uf .* H + egcno.Bf)
+        I = Flux.sigmoid_fast.(egcno.Wi .* W + egcno.Ui .* H + egcno.Bi)
+        O = Flux.sigmoid_fast.(egcno.Wo .* W + egcno.Uo .* H + egcno.Bo)
+        C̃ = Flux.tanh_fast.(egcno.Wc .* W + egcno.Uc .* H + egcno.Bc)
+        C = F .* C + I .* C̃
+        H = O .* tanh_fast.(C)
+        W = H
+        egcno.conv(tg.snapshots[i], x[i]; conv_weight = H)
+    end
+    return X
+end
  
-# function Base.show(io::IO, egcno::EvolveGCNO)
-#     print(io, "EvolveGCNO($(egcno.in) => $(egcno.out))")
-# end
+function Base.show(io::IO, egcno::EvolveGCNO)
+    print(io, "EvolveGCNO($(egcno.in) => $(egcno.out))")
+end
