def test_fullbatch_cluster_models(model_type):
G = example_graph_random(n_nodes=50)
generator = ClusterNodeGenerator(G, clusters=10)
nodes = G.nodes()[:40]
gen = generator.flow(nodes, targets=np.ones(len(nodes)))
gnn = model_type(
generator=generator,
layer_sizes=[16, 16, 1],
activations=["relu", "relu", "relu"],
)
model = tf.keras.Model(*gnn.in_out_tensors())
model.compile(optimizer="adam", loss="binary_crossentropy")
history = model.fit(gen, validation_data=gen, epochs=2)
results = model.evaluate(gen)
# this doesn't work for any cluster models including ClusterGCN
# because the model spits out predictions with shapes:
# [(1, cluster_1_size, feat_size), (1, cluster_2_size, feat_size)...]
# and attempts to concatenate along axis 0
# predictions = model.predict(gen)
x_in, x_out = gnn.in_out_tensors()
x_out_flat = tf.squeeze(x_out, 0)
embedding_model = tf.keras.Model(inputs=x_in, outputs=x_out_flat)
predictions = embedding_model.predict(gen)
assert predictions.shape == (len(nodes), 1)
def test_benchmark_ClusterGCN_generator(benchmark, q):
G = example_graph_random(feature_size=10, n_nodes=1000, n_edges=5000)
generator = ClusterNodeGenerator(G, clusters=10, q=q)
seq = generator.flow(G.nodes())
# iterate over all the batches
benchmark(lambda: list(seq))
def test_ClusterGCN_apply():
G = create_stellargraph()
generator = ClusterNodeGenerator(G)
cluster_gcn_model = ClusterGCN(
layer_sizes=[2], generator=generator, activations=["relu"], dropout=0.0
)
x_in, x_out = cluster_gcn_model.build()
model = keras.Model(inputs=x_in, outputs=x_out)
# Check fit_generator method
preds_2 = model.predict_generator(generator.flow(["a", "b", "c"]))
assert preds_2.shape == (1, 3, 2)
def test_ClusterGCN_apply():
G, _ = create_graph_features()
generator = ClusterNodeGenerator(G)
cluster_gcn_model = ClusterGCN(layer_sizes=[2],
generator=generator,
activations=["relu"],
dropout=0.0)
x_in, x_out = cluster_gcn_model.in_out_tensors()
model = keras.Model(inputs=x_in, outputs=x_out)
# Check fit method
preds_2 = model.predict(generator.flow(["a", "b", "c"]))
assert preds_2.shape == (1, 3, 2)
def _fit_deep_graph_infomax(train_graph, params, model_name):
"""Train unsupervised Deep Graph Infomax."""
if "gcn_dgi" in model_name or "gat_dgi" in model_name:
if "cluster" in model_name:
generator = ClusterNodeGenerator(
train_graph, clusters=params["clusters"],
q=params["clusters_q"])
else:
generator = FullBatchNodeGenerator(train_graph, sparse=False)
if "gcn_dgi" in model_name:
embedding_layer = GCN(
layer_sizes=[params["embedding_dimension"]],
activations=["relu"], generator=generator)
elif "gat_dgi" in model_name:
embedding_layer = GAT(
layer_sizes=[params["embedding_dimension"]],
activations=["relu"], generator=generator, attn_heads=8)
elif model_name == "graphsage_dgi":
generator = GraphSAGENodeGenerator(
train_graph, batch_size=50, num_samples=[5])
embedding_layer = GraphSAGE(
layer_sizes=[params["embedding_dimension"]], activations=["relu"],
generator=generator
)
else:
raise ValueError(f"Unknown mode name {model_name}")
embedding_model = _execute_deep_graph_infomax(
train_graph, embedding_layer, generator, params)
# Here the models can be both inductive and transductive
if model_name in ["gcn_dgi", "gat_dgi", "graphsage_dgi"]:
return embedding_model.predict(
generator.flow(train_graph.nodes()))
else:
return embedding_model
def test_ClusterNodeSequence_cluster_without_targets():
G = create_stellargraph()
generator = ClusterNodeGenerator(G, clusters=2, q=1)
seq = generator.flow(node_ids=["a"], targets=[0])
_ = list(seq)