From f40afd26e867589be4d95f5ce5fb01f78b14c530 Mon Sep 17 00:00:00 2001 From: Anant Thazhemadam Date: Sat, 9 Oct 2021 13:27:52 +0530 Subject: [PATCH 1/6] modulify things Signed-off-by: Anant Thazhemadam --- src/AtomicGraphNets.jl | 2 +- src/layers.jl | 235 ------------------------------------- src/layers/conv/agnconv.jl | 102 ++++++++++++++++ src/layers/layers.jl | 48 ++++++++ src/layers/pool/agnpool.jl | 94 +++++++++++++++ test/layer_tests.jl | 2 +- 6 files changed, 246 insertions(+), 237 deletions(-) delete mode 100755 src/layers.jl create mode 100644 src/layers/conv/agnconv.jl create mode 100755 src/layers/layers.jl create mode 100644 src/layers/pool/agnpool.jl diff --git a/src/AtomicGraphNets.jl b/src/AtomicGraphNets.jl index 309c3da..fc53335 100755 --- a/src/AtomicGraphNets.jl +++ b/src/AtomicGraphNets.jl @@ -1,7 +1,7 @@ module AtomicGraphNets export AGNConv, AGNPool#, AGNConvDEQ -include("layers.jl") +include("layers/layers.jl") using .Layers: AGNConv, AGNPool include("models.jl") diff --git a/src/layers.jl b/src/layers.jl deleted file mode 100755 index 2e06b41..0000000 --- a/src/layers.jl +++ /dev/null @@ -1,235 +0,0 @@ -module Layers - -using Flux -using Flux: glorot_uniform, normalise, @functor#, destructure -using Zygote: @adjoint, @nograd -using LinearAlgebra, SparseArrays -using Statistics -using SimpleWeightedGraphs -using ChemistryFeaturization -#using DifferentialEquations, DiffEqSensitivity - -""" - AGNConv(in=>out) - -Atomic graph convolutional layer. Almost identical to GCNConv from GeometricFlux but adapted to be most similar to Tian's original AGNN structure, so explicitly has self and convolutional weights separately. - -# Fields -- `selfweight::Array{T,2}`: weights applied to features at a node -- `convweight::Array{T,2}`: convolutional weights -- `bias::Array{T,2}`: additive bias (second dimension is always 1 because only learnable per-feature, not per-node) -- `σ::F`: activation function (will be applied before `reg_norm` to outputs), defaults to softplus - -# Arguments -- `in::Integer`: the dimension of input features. -- `out::Integer`: the dimension of output features. -- `σ=softplus`: activation function -- `initW=glorot_uniform`: initialization function for weights -- `initb=zeros`: initialization function for biases - -""" -struct AGNConv{T,F} - selfweight::Array{T,2} - convweight::Array{T,2} - bias::Array{T,2} - σ::F -end - -function AGNConv( - ch::Pair{<:Integer,<:Integer}, - σ = softplus; - initW = glorot_uniform, - initb = zeros, - T::DataType = Float64, -) - selfweight = T.(initW(ch[2], ch[1])) - convweight = T.(initW(ch[2], ch[1])) - b = T.(initb(ch[2], 1)) - AGNConv(selfweight, convweight, b, σ) -end - -@functor AGNConv - -""" - Define action of layer on inputs: do a graph convolution, add this (weighted by convolutional weight) to the features themselves (weighted by self weight) and the per-feature bias (concatenated to match number of nodes in graph). - -# Arguments -- input: a FeaturizedAtoms object, or graph_laplacian, encoded_features - -# Note -In the case of providing two matrices, the following conditions must hold: -- `lapl` must be square and of dimension N x N where N is the number of nodes in the graph -- `X` (encoded features) must be of dimension M x N, where M is `size(l.convweight)[2]` (or equivalently, `size(l.selfweight)[2]`) -""" -function (l::AGNConv{T,F})(lapl::Matrix{<:Real}, X::Matrix{<:Real}) where {T<:Real,F} - # should we put dimension checks here? Could allow more informative errors, but would likely introduce performance penalty. For now it's just in docstring. - out_mat = - T.( - normalise( - l.σ.( - l.convweight * X * lapl + - l.selfweight * X + - reduce(hcat, l.bias for i = 1:size(X, 2)), - ), - dims = [1, 2], - ), - ) - lapl, out_mat -end - -# alternate signature so FeaturizedAtoms can be fed into first layer -(l::AGNConv)(a::FeaturizedAtoms{AtomGraph,GraphNodeFeaturization}) = - l(a.atoms.laplacian, a.encoded_features) - -# signature to splat appropriately -(l::AGNConv)(t::Tuple{Matrix{R1},Matrix{R2}}) where {R1<:Real,R2<:Real} = l(t...) - -# fixes from Dhairya so backprop works -@adjoint function SparseMatrixCSC{T,N}(arr) where {T,N} - SparseMatrixCSC{T,N}(arr), Δ -> (collect(Δ),) -end -@nograd LinearAlgebra.diagm - -@adjoint function Broadcast.broadcasted(Float32, a::SparseMatrixCSC{T,N}) where {T,N} - Float32.(a), Δ -> (nothing, T.(Δ)) -end -@nograd issymmetric - -@adjoint function Broadcast.broadcasted(Float64, a::SparseMatrixCSC{T,N}) where {T,N} - Float64.(a), Δ -> (nothing, T.(Δ)) -end - -@adjoint function softplus(x::Real) - y = softplus(x) - return y, Δ -> (Δ * σ(x),) -end - -""" -Custom pooling layer that outputs a fixed-length feature vector irrespective of input dimensions, for consistent handling of different-sized graphs feeding to fully-connected dense layers afterwards. Adapted from Flux's MeanPool. - -It accepts a pooling width and will adjust stride and/or padding such that the output vector length is correct. -""" -struct AGNPool - pool_func::Function - dim::Int64 - str::Int64 - pad::Int64 - function AGNPool( - pool_type::String, - in_num_features::Int64, - out_num_features::Int64, - pool_width_frac::Float64, - ) - @assert in_num_features >= out_num_features "I don't think you actually want to pool to a LONGER vector, do you?" - dim, str, pad = - compute_pool_params(in_num_features, out_num_features, Float64(pool_width_frac)) - if pool_type == "max" - pool_func = Flux.maxpool - elseif pool_type == "mean" - pool_func = Flux.meanpool - end - new(pool_func, dim, str, pad) - end -end - -pool_out_features(num_f::Int64, dim::Int64, stride::Int64, pad::Int64) = - Int64(floor((num_f + 2 * pad - dim) / stride + 1)) - -""" -Helper function to work out dim, pad, and stride for desired number of output features, given a fixed pooling width. -""" -function compute_pool_params( - num_f_in::Int64, - num_f_out::Int64, - dim_frac::AbstractFloat; - start_dim = Int64(round(dim_frac * num_f_in)), - start_str = Int64(floor(num_f_in / num_f_out)), -) - # take starting guesses - dim = start_dim - str = start_str - p_numer = str * (num_f_out - 1) - (num_f_in - dim) - if p_numer < 0 - p_numer == -1 ? dim = dim + 1 : str = str + 1 - end - p_numer = str * (num_f_out - 1) - (num_f_in - dim) - if p_numer < 0 - error("problem, negative p!") - end - if p_numer % 2 == 0 - pad = Int64(p_numer / 2) - else - dim = dim - 1 - pad = Int64((str * (num_f_out - 1) - (num_f_in - dim)) / 2) - end - out_fea_len = pool_out_features(num_f_in, dim, str, pad) - if !(out_fea_len == num_f_out) - print("problem, output feature wrong length!") - end - # check if pad gets comparable to width... - if pad >= 0.8 * dim - @warn "specified pooling width was hard to satisfy without nonsensically large padding relative to width, had to increase from desired width" - dim, str, pad = compute_pool_params( - num_f_in, - num_f_out, - dim_frac, - start_dim = Int64(round(1.2 * start_dim)), - ) - end - dim, str, pad -end - -function (m::AGNPool)(feat::Matrix{<:Real}) - # compute what pad and stride need to be... - x = reshape(feat, (size(feat)..., 1, 1)) - # do mean pooling across feature direction, average across all nodes in graph - # TODO: decide if this approach makes sense or if there's a smarter way - pdims = PoolDims(x, (m.dim, 1); padding = (m.pad, 0), stride = (m.str, 1)) - mean(m.pool_func(x, pdims), dims = 2)[:, :, 1, 1] -end - -# alternate signatures so it can take output directly from AGNConv layer -(m::AGNPool)(lapl::Matrix{<:Real}, out_mat::Matrix{<:Real}) = m(out_mat) -(m::AGNPool)(t::Tuple{Matrix{R1},Matrix{R2}}) where {R1<:Real,R2<:Real} = m(t[2]) - -# following commented out for now because it only runs suuuuper slowly but slows down precompilation a lot -""" -# DEQ-style model where we treat the convolution as a SteadyStateProblem -struct AGNConvDEQ{T,F} - conv::AGNConv{T,F} -end - -function AGNConvDEQ(ch::Pair{<:Integer,<:Integer}, σ=softplus; initW=glorot_uniform, initb=glorot_uniform, T::DataType=Float32, bias::Bool=true) - conv = AGNConv(ch, σ; initW=initW, initb=initb, T=T) - AGNConvDEQ(conv) -end - -@functor AGNConvDEQ - -# set up SteadyStateProblem where the derivative is the convolution operation -# (we want the "fixed point" of the convolution) -# need it in the form f(u,p,t) (but t doesn't matter) -# u is the features, p is the parameters of conv -# re(p) reconstructs the convolution with new parameters p -function (l::AGNConvDEQ)(fa::FeaturizedAtoms) - p,re = Flux.destructure(l.conv) - # do one convolution to get initial guess - guess = l.conv(gr)[2] - - f = function (dfeat,feat,p,t) - input = gr - input.encoded_features = reshape(feat,size(guess)) - output = re(p)(input) - dfeat .= vec(output[2]) .- vec(input.encoded_features) - end - - prob = SteadyStateProblem{true}(f, vec(guess), p) - #return solve(prob, DynamicSS(Tsit5())).u - alg = SSRootfind() - #alg = SSRootfind(nlsolve = (f,u0,abstol) -> (res=SteadyStateDiffEq.NLsolve.nlsolve(f,u0,autodiff=:forward,method=:anderson,iterations=Int(1e6),ftol=abstol);res.zero)) - out_mat = reshape(solve(prob, alg, sensealg = SteadyStateAdjoint(autodiff = false, autojacvec = ZygoteVJP())).u,size(guess)) - return AtomGraph(gr.graph, gr.elements, out_mat, gr.featurization) -end -""" - -end diff --git a/src/layers/conv/agnconv.jl b/src/layers/conv/agnconv.jl new file mode 100644 index 0000000..252347d --- /dev/null +++ b/src/layers/conv/agnconv.jl @@ -0,0 +1,102 @@ +using Flux +using Flux: glorot_uniform, normalise, @functor#, destructure +using Zygote: @adjoint, @nograd +using LinearAlgebra, SparseArrays +using Statistics +using ChemistryFeaturization +#using DifferentialEquations, DiffEqSensitivity + +""" + AGNConv(in=>out) + +Atomic graph convolutional layer. Almost identical to GCNConv from GeometricFlux but adapted to be most similar to Tian's original AGNN structure, so explicitly has self and convolutional weights separately. + +# Fields +- `selfweight::Array{T,2}`: weights applied to features at a node +- `convweight::Array{T,2}`: convolutional weights +- `bias::Array{T,2}`: additive bias (second dimension is always 1 because only learnable per-feature, not per-node) +- `σ::F`: activation function (will be applied before `reg_norm` to outputs), defaults to softplus + +# Arguments +- `in::Integer`: the dimension of input features. +- `out::Integer`: the dimension of output features. +- `σ=softplus`: activation function +- `initW=glorot_uniform`: initialization function for weights +- `initb=zeros`: initialization function for biases + +""" +struct AGNConv{T,F} + selfweight::Array{T,2} + convweight::Array{T,2} + bias::Array{T,2} + σ::F +end + +function AGNConv( + ch::Pair{<:Integer,<:Integer}, + σ = softplus; + initW = glorot_uniform, + initb = zeros, + T::DataType = Float64, +) + selfweight = T.(initW(ch[2], ch[1])) + convweight = T.(initW(ch[2], ch[1])) + b = T.(initb(ch[2], 1)) + AGNConv(selfweight, convweight, b, σ) +end + +@functor AGNConv + +""" + Define action of layer on inputs: do a graph convolution, add this (weighted by convolutional weight) to the features themselves (weighted by self weight) and the per-feature bias (concatenated to match number of nodes in graph). + +# Arguments +- input: a FeaturizedAtoms object, or graph_laplacian, encoded_features + +# Note +In the case of providing two matrices, the following conditions must hold: +- `lapl` must be square and of dimension N x N where N is the number of nodes in the graph +- `X` (encoded features) must be of dimension M x N, where M is `size(l.convweight)[2]` (or equivalently, `size(l.selfweight)[2]`) +""" +function (l::AGNConv{T,F})(lapl::Matrix{<:Real}, X::Matrix{<:Real}) where {T<:Real,F} + # should we put dimension checks here? Could allow more informative errors, but would likely introduce performance penalty. For now it's just in docstring. + out_mat = + T.( + normalise( + l.σ.( + l.convweight * X * lapl + + l.selfweight * X + + reduce(hcat, l.bias for i = 1:size(X, 2)), + ), + dims = [1, 2], + ), + ) + lapl, out_mat +end + +# alternate signature so FeaturizedAtoms can be fed into first layer +(l::AGNConv)(a::FeaturizedAtoms{AtomGraph,GraphNodeFeaturization}) = + l(a.atoms.laplacian, a.encoded_features) + +# signature to splat appropriately +(l::AGNConv)(t::Tuple{Matrix{R1},Matrix{R2}}) where {R1<:Real,R2<:Real} = l(t...) + +# fixes from Dhairya so backprop works +@adjoint function SparseMatrixCSC{T,N}(arr) where {T,N} + SparseMatrixCSC{T,N}(arr), Δ -> (collect(Δ),) +end +@nograd LinearAlgebra.diagm + +@adjoint function Broadcast.broadcasted(Float32, a::SparseMatrixCSC{T,N}) where {T,N} + Float32.(a), Δ -> (nothing, T.(Δ)) +end +@nograd issymmetric + +@adjoint function Broadcast.broadcasted(Float64, a::SparseMatrixCSC{T,N}) where {T,N} + Float64.(a), Δ -> (nothing, T.(Δ)) +end + +@adjoint function softplus(x::Real) + y = softplus(x) + return y, Δ -> (Δ * σ(x),) +end diff --git a/src/layers/layers.jl b/src/layers/layers.jl new file mode 100755 index 0000000..53acc2d --- /dev/null +++ b/src/layers/layers.jl @@ -0,0 +1,48 @@ +module Layers + +#using DifferentialEquations, DiffEqSensitivity + +include("conv/agnconv.jl") +include("pool/agnpool.jl") + +# following commented out for now because it only runs suuuuper slowly but slows down precompilation a lot +""" +# DEQ-style model where we treat the convolution as a SteadyStateProblem +struct AGNConvDEQ{T,F} + conv::AGNConv{T,F} +end + +function AGNConvDEQ(ch::Pair{<:Integer,<:Integer}, σ=softplus; initW=glorot_uniform, initb=glorot_uniform, T::DataType=Float32, bias::Bool=true) + conv = AGNConv(ch, σ; initW=initW, initb=initb, T=T) + AGNConvDEQ(conv) +end + +@functor AGNConvDEQ + +# set up SteadyStateProblem where the derivative is the convolution operation +# (we want the "fixed point" of the convolution) +# need it in the form f(u,p,t) (but t doesn't matter) +# u is the features, p is the parameters of conv +# re(p) reconstructs the convolution with new parameters p +function (l::AGNConvDEQ)(fa::FeaturizedAtoms) + p,re = Flux.destructure(l.conv) + # do one convolution to get initial guess + guess = l.conv(gr)[2] + + f = function (dfeat,feat,p,t) + input = gr + input.encoded_features = reshape(feat,size(guess)) + output = re(p)(input) + dfeat .= vec(output[2]) .- vec(input.encoded_features) + end + + prob = SteadyStateProblem{true}(f, vec(guess), p) + #return solve(prob, DynamicSS(Tsit5())).u + alg = SSRootfind() + #alg = SSRootfind(nlsolve = (f,u0,abstol) -> (res=SteadyStateDiffEq.NLsolve.nlsolve(f,u0,autodiff=:forward,method=:anderson,iterations=Int(1e6),ftol=abstol);res.zero)) + out_mat = reshape(solve(prob, alg, sensealg = SteadyStateAdjoint(autodiff = false, autojacvec = ZygoteVJP())).u,size(guess)) + return AtomGraph(gr.graph, gr.elements, out_mat, gr.featurization) +end +""" + +end diff --git a/src/layers/pool/agnpool.jl b/src/layers/pool/agnpool.jl new file mode 100644 index 0000000..d0f2a50 --- /dev/null +++ b/src/layers/pool/agnpool.jl @@ -0,0 +1,94 @@ +using Flux +using Flux: glorot_uniform, normalise, @functor#, destructure +using Zygote: @adjoint, @nograd +using LinearAlgebra +using ChemistryFeaturization +#using DifferentialEquations, DiffEqSensitivity + +""" +Custom pooling layer that outputs a fixed-length feature vector irrespective of input dimensions, for consistent handling of different-sized graphs feeding to fully-connected dense layers afterwards. Adapted from Flux's MeanPool. + +It accepts a pooling width and will adjust stride and/or padding such that the output vector length is correct. +""" +struct AGNPool + pool_func::Function + dim::Int64 + str::Int64 + pad::Int64 + function AGNPool( + pool_type::String, + in_num_features::Int64, + out_num_features::Int64, + pool_width_frac::Float64, + ) + @assert in_num_features >= out_num_features "I don't think you actually want to pool to a LONGER vector, do you?" + dim, str, pad = + compute_pool_params(in_num_features, out_num_features, Float64(pool_width_frac)) + if pool_type == "max" + pool_func = Flux.maxpool + elseif pool_type == "mean" + pool_func = Flux.meanpool + end + new(pool_func, dim, str, pad) + end +end + +pool_out_features(num_f::Int64, dim::Int64, stride::Int64, pad::Int64) = + Int64(floor((num_f + 2 * pad - dim) / stride + 1)) + +""" +Helper function to work out dim, pad, and stride for desired number of output features, given a fixed pooling width. +""" +function compute_pool_params( + num_f_in::Int64, + num_f_out::Int64, + dim_frac::AbstractFloat; + start_dim = Int64(round(dim_frac * num_f_in)), + start_str = Int64(floor(num_f_in / num_f_out)), +) + # take starting guesses + dim = start_dim + str = start_str + p_numer = str * (num_f_out - 1) - (num_f_in - dim) + if p_numer < 0 + p_numer == -1 ? dim = dim + 1 : str = str + 1 + end + p_numer = str * (num_f_out - 1) - (num_f_in - dim) + if p_numer < 0 + error("problem, negative p!") + end + if p_numer % 2 == 0 + pad = Int64(p_numer / 2) + else + dim = dim - 1 + pad = Int64((str * (num_f_out - 1) - (num_f_in - dim)) / 2) + end + out_fea_len = pool_out_features(num_f_in, dim, str, pad) + if !(out_fea_len == num_f_out) + print("problem, output feature wrong length!") + end + # check if pad gets comparable to width... + if pad >= 0.8 * dim + @warn "specified pooling width was hard to satisfy without nonsensically large padding relative to width, had to increase from desired width" + dim, str, pad = compute_pool_params( + num_f_in, + num_f_out, + dim_frac, + start_dim = Int64(round(1.2 * start_dim)), + ) + end + dim, str, pad +end + +function (m::AGNPool)(feat::Matrix{<:Real}) + # compute what pad and stride need to be... + x = reshape(feat, (size(feat)..., 1, 1)) + # do mean pooling across feature direction, average across all nodes in graph + # TODO: decide if this approach makes sense or if there's a smarter way + pdims = PoolDims(x, (m.dim, 1); padding = (m.pad, 0), stride = (m.str, 1)) + mean(m.pool_func(x, pdims), dims = 2)[:, :, 1, 1] +end + +# alternate signatures so it can take output directly from AGNConv layer +(m::AGNPool)(lapl::Matrix{<:Real}, out_mat::Matrix{<:Real}) = m(out_mat) +(m::AGNPool)(t::Tuple{Matrix{R1},Matrix{R2}}) where {R1<:Real,R2<:Real} = m(t[2]) diff --git a/test/layer_tests.jl b/test/layer_tests.jl index e5aef09..6e0ea0a 100755 --- a/test/layer_tests.jl +++ b/test/layer_tests.jl @@ -2,7 +2,7 @@ using Test using ChemistryFeaturization using SimpleWeightedGraphs -include("../src/layers.jl") +include("../src/layers/layers.jl") using .Layers: AGNConv, AGNPool @testset "AGNConv" begin From 3641a75e523e3f4542e86c1e583cabad354576f6 Mon Sep 17 00:00:00 2001 From: Anant Thazhemadam Date: Sat, 9 Oct 2021 15:25:10 +0530 Subject: [PATCH 2/6] start sketching out eigenpool pool layer Signed-off-by: Anant Thazhemadam --- src/layers/pool/eigenpool.jl | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 src/layers/pool/eigenpool.jl diff --git a/src/layers/pool/eigenpool.jl b/src/layers/pool/eigenpool.jl new file mode 100644 index 0000000..0d6dd90 --- /dev/null +++ b/src/layers/pool/eigenpool.jl @@ -0,0 +1,27 @@ +using Flux +using Flux: glorot_uniform, normalise, @functor#, destructure +using Zygote: @adjoint, @nograd +using LinearAlgebra +using ChemistryFeaturization + +#= Pooling Layer based on +Graph Convolutional Networks with EigenPooling - Yao Ma, Suhang Wang, Charu C. Aggarwal, Jiliang Tang +https://arxiv.org/abs/1904.13107 +=# + +# TBD - what other fields would be necessary for the pooling layer itself? +struct EigenPool + pool_func::Function + function EigenPool() + new(eigen_pool_func) + end +end + +function eigen_pool_func(adj_graph::Matrix{<:Real}, features::Matrix{<:Real}) + # graph_coarsening on adj_graph + # apply EigenPooling operator on the pooled features to generate new node representations corresponding to coarsened graph +end + +function (m::EigenPool)(adj_graph::Matrix{<:Real}, features::Matrix{<:Real}) + return m.pool_func(adj_graph, features) +end From bb512f32c6db974111321d236ffaa6642cf9d8aa Mon Sep 17 00:00:00 2001 From: Anant Thazhemadam Date: Sat, 9 Oct 2021 18:04:27 +0530 Subject: [PATCH 3/6] add constructors and separate coarsening and pooling Signed-off-by: Anant Thazhemadam --- src/layers/pool/eigenpool.jl | 32 ++++++++++++++++++++++++++------ 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/src/layers/pool/eigenpool.jl b/src/layers/pool/eigenpool.jl index 0d6dd90..c4d18c4 100644 --- a/src/layers/pool/eigenpool.jl +++ b/src/layers/pool/eigenpool.jl @@ -7,21 +7,41 @@ using ChemistryFeaturization #= Pooling Layer based on Graph Convolutional Networks with EigenPooling - Yao Ma, Suhang Wang, Charu C. Aggarwal, Jiliang Tang https://arxiv.org/abs/1904.13107 + +Since the size of the matrices we are dealing with isn't as huge as the one in the original work, +let us try using this as a global pooling mechanism instead, i.e., the adjacency matrix of the crystal +itself would be only "subgraph" from which we "coarsen" to a single node and pool accordingly, which +in theory would give us the overall graph representation =# # TBD - what other fields would be necessary for the pooling layer itself? struct EigenPool - pool_func::Function + pool::Function # pooling operator applied over the adjacency matrix function EigenPool() - new(eigen_pool_func) + new(eigen_pooling) end end -function eigen_pool_func(adj_graph::Matrix{<:Real}, features::Matrix{<:Real}) - # graph_coarsening on adj_graph - # apply EigenPooling operator on the pooled features to generate new node representations corresponding to coarsened graph +function eigen_pooling(graph::Matrix{<:Real}, features::Matrix{<:Real}) + L = Atoms.normalized_laplacian(graph) # get the laplacian matrix + L_eigen_vectors = eigvecs(L) # find eigen vectors for L + result = Vector() + + for i = 1:size(L_eigen_vectors)[1] + L_i = L_eigen_vectors[:, i] + push!(result, L_i'*features) + end + + # using an agreeable H and then return H elements of result hcatt-ed into a single 1xdH vector + end +#= + = Given an adjacency graph and the corresponding matrix of node features, return the + = pooled node features which is the final graph representation fed to the dense layer + = The adj_graph would be `FeaturizedAtoms.atoms` and features would be the same as usual + =# function (m::EigenPool)(adj_graph::Matrix{<:Real}, features::Matrix{<:Real}) - return m.pool_func(adj_graph, features) + features = m.pool(adj_graph, features) + return adj_graph, features end From cad8cc05b0a26367a6d6d856c573ab7ecc125434 Mon Sep 17 00:00:00 2001 From: Anant Thazhemadam Date: Tue, 12 Oct 2021 19:03:38 +0530 Subject: [PATCH 4/6] try to use eigenpooling as the global pooling mechanism Signed-off-by: Anant Thazhemadam --- src/layers/pool/eigenpool.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/layers/pool/eigenpool.jl b/src/layers/pool/eigenpool.jl index c4d18c4..1c32cdc 100644 --- a/src/layers/pool/eigenpool.jl +++ b/src/layers/pool/eigenpool.jl @@ -22,8 +22,9 @@ struct EigenPool end end -function eigen_pooling(graph::Matrix{<:Real}, features::Matrix{<:Real}) - L = Atoms.normalized_laplacian(graph) # get the laplacian matrix +# here, L is the laplacian matrix +# this probably needs to be optimized. +function eigen_pooling(L::Matrix{<:Real}, features::Matrix{<:Real}) L_eigen_vectors = eigvecs(L) # find eigen vectors for L result = Vector() @@ -33,7 +34,8 @@ function eigen_pooling(graph::Matrix{<:Real}, features::Matrix{<:Real}) end # using an agreeable H and then return H elements of result hcatt-ed into a single 1xdH vector - + result = hcat(result...)' + reshape(result, length(result), 1) # return it as a dHx1 Matrix end #= From a38fdf43a9aee58a8f2fd9d172efcbd9d329111d Mon Sep 17 00:00:00 2001 From: Anant Thazhemadam Date: Wed, 13 Oct 2021 19:28:12 +0530 Subject: [PATCH 5/6] allow users to specify the EigenPool's output feature size Return an output feature of specified size such that `d * H` + `length(zero padding)` = `pooled_feature_length` --- src/layers/pool/eigenpool.jl | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/src/layers/pool/eigenpool.jl b/src/layers/pool/eigenpool.jl index 1c32cdc..0c81965 100644 --- a/src/layers/pool/eigenpool.jl +++ b/src/layers/pool/eigenpool.jl @@ -17,24 +17,34 @@ in theory would give us the overall graph representation # TBD - what other fields would be necessary for the pooling layer itself? struct EigenPool pool::Function # pooling operator applied over the adjacency matrix - function EigenPool() - new(eigen_pooling) + out_feature_size::Int64 + function EigenPool(out_feature_size::Int64) + new(eigen_pooling, out_feature_size) end end # here, L is the laplacian matrix # this probably needs to be optimized. -function eigen_pooling(L::Matrix{<:Real}, features::Matrix{<:Real}) +function eigen_pooling(L::Matrix{<:Real}, features::Matrix{<:Real}, out_feature_size::Int64) L_eigen_vectors = eigvecs(L) # find eigen vectors for L + + N = size(L_eigen_vectors)[1] # graph size + d = Integer(length(features)/N) # dimensions of each feature vector + + H = Integer(floor(out_feature_size/d)) # number of features to be pooled + H = H > N ? N : H # if H is greater than the number of nodes, then we pool all the nodes + + pad_len = out_feature_size - (d * H) + result = Vector() - for i = 1:size(L_eigen_vectors)[1] + for i = 1:H L_i = L_eigen_vectors[:, i] push!(result, L_i'*features) end - # using an agreeable H and then return H elements of result hcatt-ed into a single 1xdH vector - result = hcat(result...)' + # return H features + zero padded elements hcatt-ed into a single 1xdH matrix + result = hcat(result..., zeros(Float64, pad_len, 1)...)' reshape(result, length(result), 1) # return it as a dHx1 Matrix end @@ -43,7 +53,4 @@ end = pooled node features which is the final graph representation fed to the dense layer = The adj_graph would be `FeaturizedAtoms.atoms` and features would be the same as usual =# -function (m::EigenPool)(adj_graph::Matrix{<:Real}, features::Matrix{<:Real}) - features = m.pool(adj_graph, features) - return adj_graph, features -end +(m::egnpl)(adj_graph::Matrix{<:Real}, features::Matrix{<:Real}) = egnpl.pool(adj_graph, features, m.out_feature_size) From 77024f06bdc35d56375ad26eb3d3104133ba6679 Mon Sep 17 00:00:00 2001 From: Anant Thazhemadam Date: Wed, 13 Oct 2021 19:40:23 +0530 Subject: [PATCH 6/6] don't make Layer as a module Signed-off-by: Anant Thazhemadam --- src/AtomicGraphNets.jl | 2 +- src/layers/layers.jl | 4 ---- test/layer_tests.jl | 1 - 3 files changed, 1 insertion(+), 6 deletions(-) diff --git a/src/AtomicGraphNets.jl b/src/AtomicGraphNets.jl index fc53335..95649ae 100755 --- a/src/AtomicGraphNets.jl +++ b/src/AtomicGraphNets.jl @@ -2,7 +2,7 @@ module AtomicGraphNets export AGNConv, AGNPool#, AGNConvDEQ include("layers/layers.jl") -using .Layers: AGNConv, AGNPool +export AGNConv, AGNPool include("models.jl") export build_CGCNN, build_SGCNN diff --git a/src/layers/layers.jl b/src/layers/layers.jl index 53acc2d..fd421cf 100755 --- a/src/layers/layers.jl +++ b/src/layers/layers.jl @@ -1,5 +1,3 @@ -module Layers - #using DifferentialEquations, DiffEqSensitivity include("conv/agnconv.jl") @@ -44,5 +42,3 @@ function (l::AGNConvDEQ)(fa::FeaturizedAtoms) return AtomGraph(gr.graph, gr.elements, out_mat, gr.featurization) end """ - -end diff --git a/test/layer_tests.jl b/test/layer_tests.jl index 6e0ea0a..9c49d9a 100755 --- a/test/layer_tests.jl +++ b/test/layer_tests.jl @@ -3,7 +3,6 @@ using ChemistryFeaturization using SimpleWeightedGraphs include("../src/layers/layers.jl") -using .Layers: AGNConv, AGNPool @testset "AGNConv" begin # create simple line graph, populate it with feature of all ones