diff --git a/GNNLux/docs/make.jl b/GNNLux/docs/make.jl index 57df3674a..6363e6cfb 100644 --- a/GNNLux/docs/make.jl +++ b/GNNLux/docs/make.jl @@ -23,7 +23,8 @@ makedocs(; pages = ["Home" => "index.md", "API Reference" => [ "Basic" => "api/basic.md", - "Convolutional layers" => "api/conv.md"]] + "Convolutional layers" => "api/conv.md", + "Temporal Convolutional layers" => "api/temporalconv.md",]] ) diff --git a/GNNLux/docs/src/api/basic.md b/GNNLux/docs/src/api/basic.md index aef66c50a..18f943f5e 100644 --- a/GNNLux/docs/src/api/basic.md +++ b/GNNLux/docs/src/api/basic.md @@ -2,7 +2,8 @@ CurrentModule = GNNLux ``` -## GNNLayer +## Basic Layers + ```@docs GNNLayer GNNChain diff --git a/GNNLux/docs/src/api/temporalconv.md b/GNNLux/docs/src/api/temporalconv.md new file mode 100644 index 000000000..f6d565e8a --- /dev/null +++ b/GNNLux/docs/src/api/temporalconv.md @@ -0,0 +1,15 @@ +```@meta +CurrentModule = GNNLux +``` + +# Temporal Graph-Convolutional Layers + +Convolutions for time-varying graphs (temporal graphs) such as the [`TemporalSnapshotsGNNGraph`](@ref). + +## Docs + +```@autodocs +Modules = [GNNLux] +Pages = ["layers/temporalconv.jl"] +Private = false +``` diff --git a/GNNLux/src/layers/temporalconv.jl b/GNNLux/src/layers/temporalconv.jl index 64ce9b78a..e58901774 100644 --- a/GNNLux/src/layers/temporalconv.jl +++ b/GNNLux/src/layers/temporalconv.jl @@ -56,8 +56,96 @@ function Base.show(io::IO, tgcn::TGCNCell) print(io, "TGCNCell($(tgcn.in_dims) => $(tgcn.out_dims))") end +""" + TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true) + +Temporal Graph Convolutional Network (T-GCN) recurrent layer from the paper [T-GCN: A Temporal Graph Convolutional Network for Traffic Prediction](https://arxiv.org/pdf/1811.05320.pdf). + +Performs a layer of GCNConv to model spatial dependencies, followed by a Gated Recurrent Unit (GRU) cell to model temporal dependencies. + +# Arguments + +- `in`: Number of input features. +- `out`: Number of output features. +- `use_bias`: Add learnable bias. Default `true`. +- `init_weight`: Weights' initializer. Default `glorot_uniform`. +- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`. +- `init_bias`: Bias initializer. Default `zeros32`. +- `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`. +- `use_edge_weight`: If `true`, consider the edge weights in the input graph (if available). + If `add_self_loops=true` the new weights will be set to 1. + This option is ignored if the `edge_weight` is explicitly provided in the forward pass. + Default `false`. + + + +# Examples + +```julia +using GNNLux, Lux, Random + +# initialize random number generator +rng = Random.default_rng() + +# create data +g = rand_graph(rng, 5, 10) +x = rand(rng, Float32, 2, 5) + +# create TGCN layer +tgcn = TGCN(2 => 6) + +# setup layer +ps, st = LuxCore.setup(rng, tgcn) + +# forward pass +y, st = tgcn(g, x, ps, st) # result size (6, 5) +``` +""" TGCN(ch::Pair{Int, Int}; kwargs...) = GNNLux.StatefulRecurrentCell(TGCNCell(ch; kwargs...)) +""" + A3TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true) + +Attention Temporal Graph Convolutional Network (A3T-GCN) model from the paper [A3T-GCN: Attention Temporal Graph +Convolutional Network for Traffic Forecasting](https://arxiv.org/pdf/2006.11583.pdf). + +Performs a TGCN layer, followed by a soft attention layer. + +# Arguments + +- `in`: Number of input features. +- `out`: Number of output features. +- `use_bias`: Add learnable bias. Default `true`. +- `init_weight`: Weights' initializer. Default `glorot_uniform`. +- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`. +- `init_bias`: Bias initializer. Default `zeros32`. +- `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`. +- `use_edge_weight`: If `true`, consider the edge weights in the input graph (if available). + If `add_self_loops=true` the new weights will be set to 1. + This option is ignored if the `edge_weight` is explicitly provided in the forward pass. + Default `false`. +# Examples + +```julia +using GNNLux, Lux, Random + +# initialize random number generator +rng = Random.default_rng() + +# create data +g = rand_graph(rng, 5, 10) +x = rand(rng, Float32, 2, 5) + +# create A3TGCN layer +l = A3TGCN(2 => 6) + +# setup layer +ps, st = LuxCore.setup(rng, l) + +# forward pass +y, st = l(g, x, ps, st) # result size (6, 5) +``` +""" @concrete struct A3TGCN <: GNNContainerLayer{(:tgcn, :dense1, :dense2)} in_dims::Int out_dims::Int @@ -94,7 +182,7 @@ function Base.show(io::IO, l::A3TGCN) print(io, "A3TGCN($(l.in_dims) => $(l.out_dims))") end -@concrete struct GConvGRUCell <: GNNContainerLayer{(:conv_x_r, :conv_h_r, :conv_x_z, :conv_h_z, :conv_x_h, :conv_h_h)} + @concrete struct GConvGRUCell <: GNNContainerLayer{(:conv_x_r, :conv_h_r, :conv_x_z, :conv_h_z, :conv_x_h, :conv_h_h)} in_dims::Int out_dims::Int k::Int @@ -147,6 +235,45 @@ end LuxCore.outputsize(l::GConvGRUCell) = (l.out_dims,) +""" + GConvGRU(in => out, k; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32) + +Graph Convolutional Gated Recurrent Unit (GConvGRU) recurrent layer from the paper [Structured Sequence Modeling with Graph Convolutional Recurrent Networks](https://arxiv.org/pdf/1612.07659). + +Performs a layer of ChebConv to model spatial dependencies, followed by a Gated Recurrent Unit (GRU) cell to model temporal dependencies. + +# Arguments + +- `in`: Number of input features. +- `out`: Number of output features. +- `k`: Chebyshev polynomial order. +- `use_bias`: Add learnable bias. Default `true`. +- `init_weight`: Weights' initializer. Default `glorot_uniform`. +- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`. +- `init_bias`: Bias initializer. Default `zeros32`. + +# Examples + +```julia +using GNNLux, Lux, Random + +# initialize random number generator +rng = Random.default_rng() + +# create data +g = rand_graph(rng, 5, 10) +x = rand(rng, Float32, 2, 5) + +# create layer +l = GConvGRU(2 => 5, 2) + +# setup layer +ps, st = LuxCore.setup(rng, l) + +# forward pass +y, st = l(g, x, ps, st) # result size (5, 5) +``` +""" GConvGRU(ch::Pair{Int, Int}, k::Int; kwargs...) = GNNLux.StatefulRecurrentCell(GConvGRUCell(ch, k; kwargs...)) @concrete struct GConvLSTMCell <: GNNContainerLayer{(:conv_x_i, :conv_h_i, :dense_i, :conv_x_f, :conv_h_f, :dense_f, :conv_x_c, :conv_h_c, :dense_c, :conv_x_o, :conv_h_o, :dense_o)} @@ -230,6 +357,45 @@ end LuxCore.outputsize(l::GConvLSTMCell) = (l.out_dims,) +""" + GConvLSTM(in => out, k; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32) + +Graph Convolutional Long Short-Term Memory (GConvLSTM) recurrent layer from the paper [Structured Sequence Modeling with Graph Convolutional Recurrent Networks](https://arxiv.org/pdf/1612.07659). + +Performs a layer of ChebConv to model spatial dependencies, followed by a Long Short-Term Memory (LSTM) cell to model temporal dependencies. + +# Arguments + +- `in`: Number of input features. +- `out`: Number of output features. +- `k`: Chebyshev polynomial order. +- `use_bias`: Add learnable bias. Default `true`. +- `init_weight`: Weights' initializer. Default `glorot_uniform`. +- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`. +- `init_bias`: Bias initializer. Default `zeros32`. + +# Examples + +```julia +using GNNLux, Lux, Random + +# initialize random number generator +rng = Random.default_rng() + +# create data +g = rand_graph(rng, 5, 10) +x = rand(rng, Float32, 2, 5) + +# create GConvLSTM layer +l = GConvLSTM(2 => 5, 2) + +# setup layer +ps, st = LuxCore.setup(rng, l) + +# forward pass +y, st = l(g, x, ps, st) # result size (5, 5) +``` +""" GConvLSTM(ch::Pair{Int, Int}, k::Int; kwargs...) = GNNLux.StatefulRecurrentCell(GConvLSTMCell(ch, k; kwargs...)) @concrete struct DCGRUCell <: GNNContainerLayer{(:dconv_u, :dconv_r, :dconv_c)} @@ -272,8 +438,86 @@ end LuxCore.outputsize(l::DCGRUCell) = (l.out_dims,) +""" + DCGRU(in => out, k; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32) + +Diffusion Convolutional Recurrent Neural Network (DCGRU) layer from the paper [Diffusion Convolutional Recurrent Neural +Network: Data-driven Traffic Forecasting](https://arxiv.org/pdf/1707.01926). + +Performs a Diffusion Convolutional layer to model spatial dependencies, followed by a Gated Recurrent Unit (GRU) cell to model temporal dependencies. + +# Arguments + +- `in`: Number of input features. +- `out`: Number of output features. +- `k`: Diffusion step. +- `use_bias`: Add learnable bias. Default `true`. +- `init_weight`: Weights' initializer. Default `glorot_uniform`. +- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`. +- `init_bias`: Bias initializer. Default `zeros32`. + +# Examples + +```julia +using GNNLux, Lux, Random + +# initialize random number generator +rng = Random.default_rng() + +# create data +g = rand_graph(rng, 5, 10) +x = rand(rng, Float32, 2, 5) + +# create layer +l = DCGRU(2 => 5, 2) + +# setup layer +ps, st = LuxCore.setup(rng, l) + +# forward pass +y, st = l(g, x, ps, st) # result size (5, 5) +``` +""" DCGRU(ch::Pair{Int, Int}, k::Int; kwargs...) = GNNLux.StatefulRecurrentCell(DCGRUCell(ch, k; kwargs...)) +""" + EvolveGCNO(ch; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = 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. +- `use_bias`: Add learnable bias. Default `true`. +- `init_weight`: Weights' initializer. Default `glorot_uniform`. +- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`. +- `init_bias`: Bias initializer. Default `zeros32`. + +# Examples + +```julia +using GNNLux, Lux, Random + +# initialize random number generator +rng = Random.default_rng() + +# create data +tg = TemporalSnapshotsGNNGraph([rand_graph(rng, 10, 20; ndata = rand(rng, 4, 10)), rand_graph(rng, 10, 14; ndata = rand(rng, 4, 10)), rand_graph(rng, 10, 22; ndata = rand(rng, 4, 10))]) + +# create layer +l = EvolveGCNO(4 => 5) + +# setup layer +ps, st = LuxCore.setup(rng, l) + +# forward pass +y, st = l(tg, tg.ndata.x , ps, st) # result size 3, size y[1] (5, 10) +``` +""" @concrete struct EvolveGCNO <: GNNLayer in_dims::Int out_dims::Int