Adding TGCN with parameters for custom activations functions in hidden and gate activations

Author: juanmigutierrez

GraphNeuralNetworks/src/layers/temporalconv.jl

@@ -824,13 +824,16 @@ end
 
 Flux.@layer :noexpand TGCNCell
 
-function TGCNCell((in, out)::Pair{Int, Int}; kws...)
+function TGCNCell((in, out)::Pair{Int, Int}; 
+                  gate_activation = sigmoid,
+                  hidden_activation = tanh,
+                  kws...)
     conv_z = GNNChain(GCNConv(in => out, relu; kws...), GCNConv(out => out; kws...))
-    dense_z = Dense(2*out => out, sigmoid)
+    dense_z = Dense(2*out => out, gate_activation)
     conv_r = GNNChain(GCNConv(in => out, relu; kws...), GCNConv(out => out; kws...))
-    dense_r = Dense(2*out => out, sigmoid)
+    dense_r = Dense(2*out => out, gate_activation)
     conv_h = GNNChain(GCNConv(in => out, relu; kws...), GCNConv(out => out; kws...))
-    dense_h = Dense(2*out => out, tanh)
+    dense_h = Dense(2*out => out, hidden_activation)
     return TGCNCell(in, out, conv_z, dense_z, conv_r, dense_r, conv_h, dense_h)
 end
 
@@ -858,13 +861,18 @@ function Base.show(io::IO, cell::TGCNCell)
 end
 
 """
-    TGCN(args...; kws...)
+    TGCN(args...; gate_activation = sigmoid, hidden_activation = tanh, kws...)
 
 Construct a recurrent layer corresponding to the [`TGCNCell`](@ref) cell.
 
 The arguments are passed to the [`TGCNCell`](@ref) constructor.
 See [`GNNRecurrence`](@ref) for more details.
 
+# Additional Parameters
+
+- `gate_activation`: Activation function for the gate mechanisms. Default `sigmoid`.
+- `hidden_activation`: Activation function for the hidden state update. Default `tanh`.
+
 # Examples
 
 ```jldoctest
@@ -878,7 +886,7 @@ julia> g = rand_graph(num_nodes, num_edges);
 
 julia> x = rand(Float32, d_in, timesteps, num_nodes);
 
-julia> layer = TGCN(d_in => d_out)
+julia> layer = TGCN(d_in => d_out, hidden_activation = relu)
 GNNRecurrence(
   TGCNCell(2 => 3),                     # 126 parameters
 )                   # Total: 18 arrays, 126 parameters, 1.469 KiB.
@@ -889,5 +897,6 @@ julia> size(y) # (d_out, timesteps, num_nodes)
 (3, 5, 5)
 ```
 """
-TGCN(args...; kws...) = GNNRecurrence(TGCNCell(args...; kws...))
+TGCN(args...; gate_activation = sigmoid, hidden_activation = tanh, kws...) = 
+    GNNRecurrence(TGCNCell(args...; gate_activation, hidden_activation, kws...))
 

GraphNeuralNetworks/test/layers/temporalconv.jl

@@ -55,6 +55,28 @@ end
     test_gradients(model, g, x, rtol = RTOL_HIGH, atol = ATOL_LOW)
 end
 
+@testitem "TGCN with custom activations" setup=[TemporalConvTestModule, TestModule] begin
+    using .TemporalConvTestModule, .TestModule
+    using Flux: relu, sigmoid
+    layer_default = TGCN(in_channel => out_channel)
+    layer_custom = TGCN(in_channel => out_channel, gate_activation = relu, hidden_activation = relu)
+    x = rand(Float32, in_channel, timesteps, g.num_nodes)
+    y_default = layer_default(g, x)
+    y_custom = layer_custom(g, x)
+    @test layer_default isa GNNRecurrence
+    @test layer_custom isa GNNRecurrence
+    @test size(y_default) == size(y_custom)
+    # Outputs should be different due to different activations
+    @test y_default != y_custom
+    test_gradients(layer_custom, g, x, rtol = RTOL_HIGH)
+
+    # interplay with GNNChain
+    model = GNNChain(TGCN(in_channel => out_channel), Dense(out_channel, 1))
+    y = model(g, x)
+    @test size(y) == (1, timesteps, g.num_nodes)
+    test_gradients(model, g, x, rtol = RTOL_HIGH, atol = ATOL_LOW)
+end
+
 @testitem "GConvLSTMCell" setup=[TemporalConvTestModule, TestModule] begin
     using .TemporalConvTestModule, .TestModule
     cell = GConvLSTMCell(in_channel => out_channel, 2)