Hi,
I was trying to Bayesianize graph neural networks using NumPyro and Jraph.
Firstly, if I do not use Jraph to define the graphs of interest, then there is no real issue. One can go ahead and define graph operations with base JAX, and couple it with Flax for NN definitions, and then proceed with random_flax_module.
However, Jraph provides a nice object for dealing with standard graphs. I have copied a model from Jraph’s examples to illustrate:
class ExplicitMLP(nn.Module):
"""A flax MLP."""
features: Sequence[int]
@nn.compact
def __call__(self, inputs):
x = inputs
for i, lyr in enumerate([nn.Dense(feat) for feat in self.features]):
x = lyr(x)
if i != len(self.features) - 1:
x = nn.relu(x)
return x
def make_embed_fn(latent_size):
def embed(inputs):
return nn.Dense(latent_size)(inputs)
return embed
def make_mlp(features):
@jraph.concatenated_args
def update_fn(inputs):
return ExplicitMLP(features)(inputs)
return update_fn
class GraphNetwork(nn.Module):
"""A flax GraphNetwork."""
mlp_features: Sequence[int]
latent_size: int
@nn.compact
def __call__(self, graph):
# Add a global parameter for graph classification computation
graph = graph._replace(globals=jnp.zeros([graph.n_node.shape[0], 1]))
embedder = jraph.GraphMapFeatures(
embed_node_fn=make_embed_fn(self.latent_size),
embed_edge_fn=make_embed_fn(self.latent_size),
embed_global_fn=make_embed_fn(self.latent_size))
net = jraph.GraphNetwork(
update_node_fn=make_mlp(self.mlp_features),
update_edge_fn=make_mlp(self.mlp_features),
# The global update outputs size 2 for binary classification.
update_global_fn=make_mlp(self.mlp_features + (2,))) # pytype: disable=unsupported-operands
return net(embedder(graph))
The flax.linen module GraphNetwork is where the main aspects of the GNN are synthesized. However, instead of accepting a static tensor as an input, it accepts a graph object (this is the Jraph GraphsTuple object). This graph object has many attributes to it, that cannot be simplified to some constant array (or atleast doing so defeats the purpose). I’m struggling to figure out how to navigate the input_shape field for random_flax_module.
Would appreciate any pointers or if I’ve missed something.
For the NumPyro part, doing something like below seems to initiate sampling (albeit very slow because of my for loop), but I’m not sure if it is correct. I can pass a graph argument to random_flax_module. I don’t know if this argument is only used to initialize the flax module, or if it’s being passed everytime the model is called.
def numpyro_model(list_of_graphs, target):
# Define base flax module
module = GraphNetwork(mlp_features=[16, 16], latent_size=16)
# Register as a random flax module (parameter inference done outside)
# TODO: check if graph argument makes sense. Would it just keep the graph constant for all calls of net?
net = random_flax_module(
"Mutag_GraphNet_Model",
module,
prior = dist.StudentT(df=4.0, scale=0.1),
input_shape=None,
graph=list_of_graphs[0]['input_graph']
)
# Prediction over a list of graphs (right now it must be looped because of issues with graph size)
# TODO: need to vectorize this
preds = jnp.zeros(shape=(len(list_of_graphs), 1))
for idx, g in enumerate(list_of_graphs):
# pred_graph = net.apply(params, g['input_graph'])
pred_graph = net(g['input_graph'])
pred = jax.nn.log_softmax(pred_graph.globals)
preds = preds.at[idx].set(pred.flatten()[0])
# Compute likelihood function
# TODO: make this the standard classification target (predict probabilities and observe binary responses)
return numpyro.sample("target", dist.Normal(jnp.exp(preds), 0.03), obs=target)
# return preds
# Initialize model
model2 = GraphNetwork(mlp_features=[16, 16], latent_size=16)
key = jax.random.PRNGKey(0)
params = model2.init(key, train_mutag_ds[0]['input_graph'])
print(parameter_overview.get_parameter_overview(params)) # uses clu (common loop utils)
del model2
# Initialize MCMC
kernel = NUTS(numpyro_model,
init_strategy = init_to_feasible(),
target_accept_prob=0.80,
max_tree_depth=10,
)
mcmc = MCMC(
kernel,
num_warmup=1000,
num_samples=1000,
num_chains=1,
progress_bar=True, # TOGGLE this...
chain_method="vectorized",
# jit_model_args=True,
)
# Run MCMC
mcmc.run(jax.random.PRNGKey(3), train_mutag_ds, y_train)
yes, arguments in random_flax_module are only used to initialize/substitute required parameters. your code looks correct to me
1 Like
Great, thanks for checking @fehiepsi! It is pretty nice that a more complicated set of objects can be handled with this functionality!
@fehiepsi do you have any thoughts on how I might go about speeding up the forward-pass here? Particularly this bit of code:
preds = jnp.zeros(shape=(len(list_of_graphs), 1))
for idx, g in enumerate(list_of_graphs):
# pred_graph = net.apply(params, g['input_graph'])
pred_graph = net(g['input_graph'])
pred = jax.nn.log_softmax(pred_graph.globals)
preds = preds.at[idx].set(pred.flatten()[0])
The list_of_graphs contains a set of Jraph GraphTuples objects (a list of inputs), where the attribute input_graph needs to be accessed for the forward pass. Using the standard jax.lax.scan or jax.lax.for_iloop constructs don’t seem to work (or atleast I can’t seem to get them to work). This chunk of code seems to be non-jittable. Would appreciate any pointers!
I think you can try batching instead of list of graphs. There are some helpers here Jraph API — Jraph 0.0.1.dev documentation
Completely missed that. Seems to work now. Thanks so much for the pointer!
For future reference, linking to the final implementation of GNNs using Jraph, Flax and NumPyro. Thanks again for the help @fehiepsi!
1 Like