Loading saved GeneralGNN

[MarcosZ]

Hello,

I'm trying to train a GeneralGNN and then export it for use at a later time. Below is the flow of my training (not using custom training, just model.fit).

model = GeneralGNN(dataset.n_labels, activation='linear', pool=None, batch_norm=False, connectivity='sum', dropout=0.0, hidden=256)
model.compile(
    optimizer=Adam(learning_rate=1e-3),
    loss=MSE,
    )

# %%
cutoff = int(dataset.n_graphs*.8)

# Important to set node_level to true for multiple graphs
loader_tr = DisjointLoader(dataset[:cutoff], node_level=True)
loader_va = DisjointLoader(dataset[cutoff:], node_level=True)

# %%
model.fit(
    loader_tr.load(),
    steps_per_epoch=loader_tr.steps_per_epoch,
    validation_data=loader_va.load(),
    validation_steps=loader_va.steps_per_epoch,
    callbacks=[EarlyStopping()],
    # epochs=10,
)

# %%
# how to save
model.save('./data/spektral.keras')

# %%
# how to load
loaded_model = load_model('./data/spektral.keras')

Right now I'm getting an error because of the fact that the MessagePassing layer is expecting a SparseTensor instead of a Tensor as is being provided, apparently.
I'm not sure how to control this because the custom data set I'm using has a SciPy sparse array as the adjacency matrix, so I figured it would work.

AssertionError: Exception encountered when calling layer 'general_conv_80' (type GeneralConv).

in user code:

    File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/spektral/layers/convolutional/general_conv.py", line 137, in call  *
        x, a, _ = self.get_inputs(inputs)
    File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/spektral/layers/convolutional/message_passing.py", line 183, in get_inputs  *
        assert K.is_sparse(a), "A must be a SparseTensor"

    AssertionError: A must be a SparseTensor


Call arguments received by layer 'general_conv_80' (type GeneralConv):
  • inputs=['tf.Tensor(shape=(None, 256), dtype=float32)', 'tf.Tensor(shape=(None, None), dtype=float32)']
  • kwargs={'training': 'None'}

In the above code, I'm able to train perfectly fine. I can then predict values that are really accurate. I can save the model.
However, when I go to load it, this is the error I get (above).

I've seen a custom training loop, but the use of model.fit is just very handy. Is there anything I can do in the dataset, perhaps, to avoid this error? I converted the adjacent matrix to SparseTensor and I failed on fit instead of load_model.

[danielegrattarola]

Hi,

that's weird, are you sure that the dataset actually contains sparse adjacency matrices? You can use dataset.apply to make sure that all graphs have sparse adjacency matrices (you'd have to implement a simple function foo(graph: Graph) -> Graph to convert the adj matrix to sparse and pass it to apply).

[MarcosZ]

Hey Daniele. Thanks for getting back to me. I tried to use transforms = spektral.transforms.adj_to_sp_tensor.AdjToSpTensor() and that didn't work either.
My first configuration is using scipy sparse matrices:

dataset[0].a

<10x10 sparse array of type '<class 'numpy.float32'>'
	with 50 stored elements in Compressed Sparse Row format>

This is where I get the above error/condition where it can train using .fit() just fine and I can .predict() in my local space, but when I go to export (load, not save) it fails. Meaning, I can save, but not load. This error is mentioned in the original post.

I then tried the above transform mentioned and get this:

dataset[0].a

SparseTensor(indices=tf.Tensor(
[[0 1]
 [0 2]
 [0 3]
 [0 4]
 [0 5]
 [0 6]
 [0 7]
 [0 8]
 [0 9]
 [1 0]
...
[0.6802 0.6679 0.1508 0.7158 0.7269 0.1075 0.282  0.6262 0.6376 0.6802
 0.7379 0.519  0.0729 0.3709 0.6679 0.1508 0.1625 0.7158 0.7643 0.5276
 0.7269 0.8192 0.6991 0.2996 0.6961 0.1075 0.7379 0.1625 0.7643 0.8192
 0.2728 0.6906 0.282  0.519  0.5276 0.6991 0.2728 0.4896 0.3669 0.6262
 0.0729 0.2996 0.6906 0.4896 0.3989 0.6376 0.3709 0.6961 0.3669 0.3989], shape=(50,), dtype=float32), dense_shape=tf.Tensor([10 10], shape=(2,), dtype=int64))

After transforming, I try to use .fit() and get this:

Traceback (most recent call last):

  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/script_ops.py", line 266, in __call__
    return func(device, token, args)

  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/script_ops.py", line 144, in __call__
    outputs = self._call(device, args)

  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/script_ops.py", line 151, in _call
    ret = self._func(*args)

  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
    return func(*args, **kwargs)

  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/data/ops/from_generator_op.py", line 263, in generator_py_func
    values = next(generator_state.get_iterator(iterator_id.numpy()))

  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/spektral/data/loaders.py", line 100, in __next__
    return self.collate(nxt)

  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/spektral/data/loaders.py", line 300, in collate
    output = to_disjoint(**packed)
...

AttributeError: 'SparseTensor' object has no attribute 'ravel'

Do you think that what's done here is the equivalent of dataset.apply or do I have to implement something specifically for .fit() to use? I'm unfamiliar with what you mean about this .apply() function.
I figured if I converted the matrices in the dataset to SparseTensors it would work, but I was confused why I couldn't just use my original scipy sparse matrices.

Let me know what you think.

[MarcosZ]

I should mention that if I write a custom train function instead of .fit() for MSE, I am able to train, predict with, export and load the model.
There are some discontinuities, it appears, somewhere in the call() function for this model and the standard Keras .fit() and .load() functions.

I'm still trying to work around it, but I'm not knowledgeable enough in this field.

[danielegrattarola]

What version of Keras are you using?

[MarcosZ]

I've installed spektral~=1.3 and it seems that the dependency of keras that came with it is 2.13.1.

[MarcosZ]

This isn't quite a solution, but I found that I'm able to just save the weights with save_weights(). However, I can only use load_weights() if I instantiate a GeneralGNN and use the predict() function. If I don't use predict() prior to loading I get another error.