GeometricFlux: GCNConv but with different graphs

I’m working with GeometricFlux and want to adapt the GCNConv layer to allow for different adjacency matrices corresponding to each piece of training data. Since this now requires two separate “streams” of input data (the adjacency matrices and the set of node features corresponding to each graph), I’m wondering about any best practices for dealing with this type of input. Most of the Flux examples I’ve seen have just one vector (or matrix) as the input. Anything particular I should be keeping in mind?

Hi, Im working on a similar problem and wondering if you were able to successfully adapt the layer?

Yes. Unfortunately the solution I have is currently closed-source because of silly academia funding/IP things, but we plan to open-source it once we publish the paper, hopefully in the next couple months. In the meantime, check out the FeaturedGraph type in GeometricFlux (may be only in master branch currently, not sure); that’s what I’m using as input to the layer.

The main conceptual thing to keep in mind in adapting it (if you want to make your own in the meantime) is that the bias terms will not be able to be trainable per-node since the number of input nodes in the graph can vary, so you can only train it per-feature. If you want to know when it comes out, I’d suggest you keep an eye on this Github org, as this is where it will be published: https://github.com/aced-differentiate

@rkurchin Thanks! This FeaturedGraph type is extremely helpful for my use case. Did you have any trouble computing the gradients when using this type combined with the GCNConv layer? I seem to get the following method error:

MethodError: no method matching (::GeometricFlux.var"#adjacency_matrix#76")(::FeaturedGraph{Array{Float64,2},Adjoint{Float64,Array{Float64,2}},Array{Float64,2},Array{Float64,1}})

Yeah, we did end up having to define a custom adjoint for that but I think you may have to wait until our paper comes out to see it. Sorry our funder is so ornery about this! We definitely will open-source it eventually, just have to follow a bunch of rules first…

@rkurchin Totally understand, thanks! :slight_smile:

1 Like

Psst…took a bit longer than anticipated, but this is finally open: https://github.com/aced-differentiate/AtomicGraphNets.jl

We ended up ditching the GeometricFlux dependency entirely and building our own graph type, which is implemented here as AtomGraph: https://github.com/aced-differentiate/ChemistryFeaturization.jl

The adjoints I mentioned earlier are here: https://github.com/aced-differentiate/AtomicGraphNets.jl/blob/58cc531a331a2132f05f499d14ad8d1e5c0822c1/src/layers.jl#L65

2 Likes

Why’d you ditch geometric flux ?

The FeaturedGraph is currently available to put graph and features together.
If you want graph to be variable, remember to turn off the cache mechanism.

GCNConv(num_features=>hidden, relu, cache=false)

I think that would help.

Any further questions, just tag me.

The FeaturedGraph type didn’t have all the checks I wanted (e.g. making sure the dimension of the feature matrix made sense given the dimensions of the graph), and because I’m using this for convolutions, I wanted to store the graph laplacian as well to avoid having to compute it every time. For those reasons and a couple others (mostly related to storing some more specific stuff related to how the featurization is done), I decided to define my own graph type, and then since I wasn’t using the GCNConv layers or the FeaturedGraph type, the GeometricFlux dependency wasn’t really necessary.

@yuehhua, I’d certainly encourage you to take a look at the AtomGraph type I defined in ChemistryFeaturization and think about incorporating some of those checks on feature matrix dimension and potentially also storing the laplacian to speed up convolution…

Yeah, I have considered to add dimension check between feature matrix and graph. In the other side, I want it to be more general as possible and they maybe don’t consistent with each other during computation. I consider a better way to add dimension check.

As for Laplacian, I still very likely to add this feature but I have some concern about how to put it in FeaturedGraph and have a consistent layout. I want it not to have one copy of matrix to reduce memory allocations.

I plan to add dimensional check first.