def _convert_to_sparse_tensor(x):
if x.ndim == 2:
return sp_matrix_to_sp_tensor(x)
elif x.ndim == 3:
s1_, s2_, s3_ = x.shape
return ops.reshape(sp_matrix_to_sp_tensor(x.reshape(s1_ * s2_, s3_)),
(s1_, s2_, s3_))
def test_modes_ops():
ns = [10, 5, 8]
x_list = [np.random.randn(n, 3) for n in ns]
a_list = [sp.csr_matrix(np.ones((n, n))) for n in ns]
x, a, i = to_disjoint(x_list, a_list)
a = sp_matrix_to_sp_tensor(a)
expected_x = np.zeros((len(ns), max(ns), 3))
for i_, n in enumerate(ns):
expected_x[i_, :n] = x_list[i_]
expected_a = np.zeros((len(ns), max(ns), max(ns)))
for i_, n in enumerate(ns):
expected_a[i_, :n, :n] = a_list[i_].A
# Disjoint signal to batch
result = ops.disjoint_signal_to_batch(x, i).numpy()
assert expected_x.shape == result.shape
assert np.allclose(expected_x, result, atol=tol)
# Disjoint adjacency to batch
result = ops.disjoint_adjacency_to_batch(a, i).numpy()
assert expected_a.shape == result.shape
assert np.allclose(expected_a, result, atol=tol)
def test_modes_ops():
X = np.array([[1, 0], [0, 1], [1, 1], [0, 0], [1, 2]])
I = np.array([0, 0, 0, 1, 1])
A_data = [1, 1, 1, 1, 1]
A_row = [0, 1, 2, 3, 4]
A_col = [1, 0, 1, 4, 3]
A_sparse = sp.csr_matrix((A_data, (A_row, A_col)), shape=(5, 5))
A_sparse_tensor = sp_matrix_to_sp_tensor(A_sparse)
# Disjoint signal to batch
expected_result = np.array([[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]],
[[0.0, 0.0], [1.0, 2.0], [0.0, 0.0]]])
result = ops.disjoint_signal_to_batch(X, I).numpy()
assert expected_result.shape == result.shape
assert np.allclose(expected_result, result, atol=tol)
# Disjoint adjacency to batch
expected_result = np.array([
[[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
[[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
])
result = ops.disjoint_adjacency_to_batch(A_sparse_tensor, I).numpy()
assert expected_result.shape == result.shape
assert np.allclose(expected_result, result, atol=tol)
def _get_input_from_dtypes(dtypes, sparse=False):
assert len(dtypes) >= 2
x = np.ones((3, 1), dtype=dtypes[0])
a = np.ones((3, 3), dtype=dtypes[1])
if sparse:
a = sp_matrix_to_sp_tensor(a)
output = [x, a]
if len(dtypes) > 2:
e = np.ones((3 * 3, 1), dtype=dtypes[2])
output.append(e)
return output
def test_sparse_model_sizes():
"""
This is a sanity check to make sure we have the same number of operations that we intend to have
"""
N = 5
F = 4
S = 3
X_in = Input(shape=(F, ), name="X_in")
A_in = Input(shape=(None, ), name="A_in", sparse=True)
E_in = Input(shape=(S, ), name="E_in")
x = np.ones(shape=(N, F))
a = np.ones(shape=(N, N))
a = sp_matrix_to_sp_tensor(a)
e = np.ones(shape=(N * N, S))
def assert_n_params(inp, out, expected_size):
model = Model(inputs=inp, outputs=out)
model.compile(optimizer="adam", loss="mean_squared_error")
print(model.count_params())
assert model.count_params() == expected_size
# for test coverage:
model([x, a, e])
X, E = XENetConv([5], 10, 20, False)([X_in, A_in, E_in])
assert_n_params([X_in, A_in, E_in], [X, E], 350)
# int vs list: 5 vs [5]
X, E = XENetConv(5, 10, 20, False)([X_in, A_in, E_in])
assert_n_params([X_in, A_in, E_in], [X, E], 350)
# t = (4+4+3+3+1)*5 = 75 # Stack Conv
# x = (4+5+5+1)*10 = 150 # Node reduce
# e = (5+1)*20 = 120 # Edge reduce
# p = 5 # Prelu
# total = t+x+e+p = 350
X, E = XENetConv(5, 10, 20, True)([X_in, A_in, E_in])
assert_n_params([X_in, A_in, E_in], [X, E], 362)
# t = (4+4+3+3+1)*5 = 75
# a = (5+1)*1 *2 = 12 # Attention
# x = (4+5+5+1)*10 = 150
# e = (5+1)*20 = 120
# p = 5 # Prelu
# total = t+x+e+p = 362
X, E = XENetConv([50, 5], 10, 20, True)([X_in, A_in, E_in])
assert_n_params([X_in, A_in, E_in], [X, E], 1292)
def _test_disjoint_mode(layer, sparse=False, **kwargs):
A = sp.block_diag(
[np.ones((N1, N1)), np.ones((N2, N2)), np.ones((N3, N3))]
).todense()
X = np.random.normal(size=(N, F))
I = np.array([0] * N1 + [1] * N2 + [2] * N3).astype(int)
A_in = Input(shape=(None,), sparse=sparse)
X_in = Input(shape=(F,))
I_in = Input(shape=(), dtype=tf.int32)
inputs = [X_in, A_in, I_in]
if sparse:
input_data = [X, sp_matrix_to_sp_tensor(A), I]
else:
input_data = [X, A, I]
layer_instance = layer(**kwargs)
output = layer_instance(inputs)
model = Model(inputs, output)
output = model(input_data)
X_pool, A_pool, I_pool, s = output
if "ratio" in kwargs.keys():
N_pool_expected = int(
np.ceil(kwargs["ratio"] * N1)
+ np.ceil(kwargs["ratio"] * N2)
+ np.ceil(kwargs["ratio"] * N3)
)
elif "k" in kwargs.keys():
N_pool_expected = int(kwargs["k"])
else:
N_pool_expected = None
N_pool_true = A_pool.shape[0]
if N_pool_expected is not None:
_check_number_of_nodes(N_pool_expected, N_pool_true)
assert X_pool.shape == (N_pool_expected, F)
assert A_pool.shape == (N_pool_expected, N_pool_expected)
assert I_pool.shape == (N_pool_expected,)
output_shape = [o.shape for o in output]
_check_output_and_model_output_shapes(output_shape, model.output_shape)
def _test_single_mode(layer, sparse=False, edges=False, **kwargs):
A_in = Input(shape=(None, ), sparse=sparse)
X_in = Input(shape=(F, ))
inputs = [X_in, A_in]
if sparse:
input_data = [X, sp_matrix_to_sp_tensor(A)]
else:
input_data = [X, A]
if edges:
E_in = Input(shape=(S, ))
inputs.append(E_in)
input_data.append(E_single)
layer_instance = layer(**kwargs)
output = layer_instance(inputs)
model = Model(inputs, output)
output = model(input_data)
assert output.shape == (N, kwargs["channels"])
def k_hop_sparse_subgraph(a, node_idx, k, transformer=None):
"""
Computes the subgraph containing all the neighbors of `node_idx` up to the k-th order.
If `a` is not the binary adjacency matrix a `transformer` should be passed.
**Arguments**
- `a`: sparse `(n_nodes, n_nodes)` graph tensor;
- `node_idx`: center node;
- `k`: order of neighbor;
- `transformer`: one of the functions from the `spektral.transforms` module,
needed to convert the binary adjacency matrix into the correct format for the model;
"""
if a.dtype != tf.float32:
a = tf.cast(a, tf.float32)
if transformer:
a = binary_adj_converter(a)
power_a = tf.sparse.eye(a.shape[0])
k_neighs = np.zeros(a.shape[0]).astype("float32").reshape(1, -1)
k_neighs[0, node_idx] = 1
for _ in range(k - 1):
power_a = dot(power_a, a)
temp = tf.sparse.slice(power_a,
start=[node_idx, 0],
size=[1, power_a.shape[0]])
k_neighs += tf.sparse.to_dense(temp)
comp_graph = tf.sparse.add(a * tf.reshape(k_neighs, (-1, 1)), a * k_neighs)
is_nonzero = tf.not_equal(comp_graph.values, 0)
comp_graph = tf.sparse.retain(comp_graph, is_nonzero)
comp_graph = tf.sign(comp_graph)
if transformer:
comp_graph = sp_tensor_to_sp_matrix(comp_graph)
comp_graph = transformer(comp_graph)
return sp_matrix_to_sp_tensor(comp_graph)
else:
return comp_graph
def _test_single_mode(layer, sparse=False, **kwargs):
A = np.ones((N, N))
X = np.random.normal(size=(N, F))
A_in = Input(shape=(None,), sparse=sparse)
X_in = Input(shape=(F,))
inputs = [X_in, A_in]
if sparse:
input_data = [X, sp_matrix_to_sp_tensor(A)]
else:
input_data = [X, A]
layer_instance = layer(**kwargs)
output = layer_instance(inputs)
model = Model(inputs, output)
output = model(input_data)
X_pool, A_pool, mask = output
if "ratio" in kwargs.keys():
N_exp = kwargs["ratio"] * N
N_pool_expected = int(np.ceil(N_exp))
elif "k" in kwargs.keys():
N_pool_expected = int(kwargs["k"])
else:
N_pool_expected = None
N_pool_true = A_pool.shape[-1]
assert N_pool_true > 0
if N_pool_expected is not None:
_check_number_of_nodes(N_pool_expected, N_pool_true)
assert X_pool.shape == (N_pool_expected, F)
assert A_pool.shape == (N_pool_expected, N_pool_expected)
output_shape = [o.shape for o in output]
_check_output_and_model_output_shapes(output_shape, model.output_shape)
def test_disjoint_2_batch():
X = np.array([[1, 0], [0, 1], [1, 1], [0, 0], [1, 2]])
I = np.array([0, 0, 0, 1, 1])
A_data = [1, 1, 1, 1, 1]
A_row = [0, 1, 2, 3, 4]
A_col = [1, 0, 1, 4, 3]
A = sp_matrix_to_sp_tensor(sp.csr_matrix((A_data, (A_row, A_col)), shape=(5, 5)))
expected_X = np.array(
[[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [1.0, 2.0], [0.0, 0.0]]]
)
expected_A = np.array(
[
[[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
[[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
]
)
result_X, result_A = layers.Disjoint2Batch()((X, A, I))
assert np.allclose(result_A, expected_A, atol=tol)
assert np.allclose(result_X, expected_X, atol=tol)
_test_get_config(layers.Disjoint2Batch)
def _test_mixed_mode(layer, sparse=False, edges=False, **kwargs):
X_batch = np.stack([X] * batch_size)
A_in = Input(shape=(N, ), sparse=sparse)
X_in = Input(shape=(N, F))
inputs = [X_in, A_in]
if sparse:
input_data = [X_batch, sp_matrix_to_sp_tensor(A)]
else:
input_data = [X_batch, A]
if edges:
E_in = Input(shape=(N * N, S))
inputs.append(E_in)
E_batch = np.stack([E_single] * batch_size)
input_data.append(E_batch)
layer_instance = layer(**kwargs)
output = layer_instance(inputs)
model = Model(inputs, output)
output = model(input_data)
assert output.shape == (batch_size, N, kwargs["channels"])
gradients = tape.gradient(loss, model.trainable_variables)
opt.apply_gradients(zip(gradients, model.trainable_variables))
return model.losses[0], model.losses[1], S_pool
np.random.seed(1)
epochs = 5000 # Training iterations
lr = 5e-4 # Learning rate
################################################################################
# LOAD DATASET
################################################################################
dataset = Cora()
adj, x, y = dataset[0].a, dataset[0].x, dataset[0].y
a_norm = normalized_adjacency(adj)
a_norm = sp_matrix_to_sp_tensor(a_norm)
F = dataset.n_node_features
y = np.argmax(y, axis=-1)
n_clusters = y.max() + 1
################################################################################
# MODEL
################################################################################
x_in = Input(shape=(F, ), name="X_in")
a_in = Input(shape=(None, ), name="A_in", sparse=True)
x_1 = GCSConv(16, activation="elu")([x_in, a_in])
x_1, a_1, s_1 = MinCutPool(n_clusters, return_selection=True)([x_1, a_in])
model = Model([x_in, a_in], [x_1, s_1])
from spektral.layers import GCNConv, GlobalSumPool
from spektral.utils.sparse import sp_matrix_to_sp_tensor
# Parameters
batch_size = 32 # Batch size
epochs = 1000 # Number of training epochs
patience = 10 # Patience for early stopping
l2_reg = 5e-4 # Regularization rate for l2
# Load data
data = MNIST()
# The adjacency matrix is stored as an attribute of the dataset.
# Create filter for GCN and convert to sparse tensor.
data.a = GCNConv.preprocess(data.a)
data.a = sp_matrix_to_sp_tensor(data.a)
# Train/valid/test split
data_tr, data_te = data[:-10000], data[-10000:]
np.random.shuffle(data_tr)
data_tr, data_va = data_tr[:-10000], data_tr[-10000:]
# We use a MixedLoader since the dataset is in mixed mode
loader_tr = MixedLoader(data_tr, batch_size=batch_size, epochs=epochs)
loader_va = MixedLoader(data_va, batch_size=batch_size)
loader_te = MixedLoader(data_te, batch_size=batch_size)
# Build model
class Net(Model):
def __init__(self, **kwargs):
def sp_matrices_to_sp_tensors(inputs):
inputs = list(inputs)
for i in range(len(inputs)):
if sp.issparse(inputs[i]):
inputs[i] = sp_matrix_to_sp_tensor(inputs[i])
return tuple(inputs)
def __call__(self, graph):
if graph.a is not None:
graph.a = sp_matrix_to_sp_tensor(graph.a)
return graph
