Spektral using on multi gpus - distributed training

[wacco]

Hello guys,
Is it possible using Spektral on multi gpus system?
I working on GNN for node level prediction using tensorflow and spektral. I using lot of data, witch i cant fit into one gpu device memory (like 20GB one big rgaph).
Do somebody some tips how i should proceed? My first idea was using distributed training with tensorflow, for more memory and more speed. But when i use tf.distribute.MirroredStrategy() i have got error like this:

tensorflow.python.framework.errors_impl.InvalidArgumentError: 3 root error(s) found.
  (0) Invalid argument:  Cannot multiply A and B because inner dimension does not match: 7230 vs. 2410.  Did you forget a transpose?  Dimensions of A: [2410, 7230).  Dimensions of B: [2410,326]
	 [[node model/cheb_conv/SparseTensorDenseMatMul/SparseTensorDenseMatMul (defined at private/var/python/model_runner/.venv/lib/python3.8/site-packages/spektral/layers/ops/matmul.py:35) ]]
	 [[Adam/ArithmeticOptimizer/AddOpsRewrite_AddN_3/_136]]
  (1) Invalid argument:  Cannot multiply A and B because inner dimension does not match: 7230 vs. 2410.  Did you forget a transpose?  Dimensions of A: [2410, 7230).  Dimensions of B: [2410,326]
	 [[node model/cheb_conv/SparseTensorDenseMatMul/SparseTensorDenseMatMul (defined at private/var/python/model_runner/.venv/lib/python3.8/site-packages/spektral/layers/ops/matmul.py:35) ]]
	 [[Adam/Adam/group_deps/_171]]
  (2) Invalid argument:  Cannot multiply A and B because inner dimension does not match: 7230 vs. 2410.  Did you forget a transpose?  Dimensions of A: [2410, 7230).  Dimensions of B: [2410,326]
	 [[node model/cheb_conv/SparseTensorDenseMatMul/SparseTensorDenseMatMul (defined at private/var/python/model_runner/.venv/lib/python3.8/site-packages/spektral/layers/ops/matmul.py:35) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_6460]

Errors may have originated from an input operation.
Input Source operations connected to node model/cheb_conv/SparseTensorDenseMatMul/SparseTensorDenseMatMul:
 cond_3/Identity_1 (defined at private/var/python/model_runner/incomingInvoice/create_model.py:102)	
 model/Cast (defined at usr/local/Cellar/python@3.8/3.8.11/Frameworks/Python.framework/Versions/3.8/lib/python3.8/threading.py:932)

Input Source operations connected to node model/cheb_conv/SparseTensorDenseMatMul/SparseTensorDenseMatMul:
 cond_3/Identity_1 (defined at private/var/python/model_runner/incomingInvoice/create_model.py:102)	
 model/Cast (defined at usr/local/Cellar/python@3.8/3.8.11/Frameworks/Python.framework/Versions/3.8/lib/python3.8/threading.py:932)

Input Source operations connected to node model/cheb_conv/SparseTensorDenseMatMul/SparseTensorDenseMatMul:
 cond_3/Identity_1 (defined at private/var/python/model_runner/incomingInvoice/create_model.py:102)	
 model/Cast (defined at usr/local/Cellar/python@3.8/3.8.11/Frameworks/Python.framework/Versions/3.8/lib/python3.8/threading.py:932)

Function call stack:
train_function -> train_function -> train_function

Bellow its my example model:

# Load dataset
dataset = Dataset(dataset_dir=dataset_dir, classes=CLASSES, use_cached=False)
print("graphs:", dataset.n_graphs, ", nodes:", dataset.n_nodes, "labels:", dataset.n_labels, ", features:", dataset.n_node_features)

# Use strategy
strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1", "GPU:2"])

# Parameters
channels = 16
learning_rate = 1e-2
epochs = 100
patience = 50
l2_reg = 5e-4
K = 3
dropout = 0.5
batch_size = 32
global_batch_size = batch_size * strategy.num_replicas_in_sync

# Data loader
loader_tr = DisjointLoader(dataset, batch_size=batch_size, node_level=True)

# Model definition
X = Input(shape=(dataset.n_node_features, ))
A = Input((None, ), sparse=True)
I = Input((batch_size, ))


with strategy.scope():
    # Define layers
    gcn0 = ChebConv(channels, K=K, activation='relu', kernel_regularizer=l2(l2_reg), use_bias=False)([X, A])
    gcn1 = ChebConv(channels * 2, K=K, activation='relu', kernel_regularizer=l2(l2_reg), use_bias=False)([gcn0, A])
    gcn2 = ChebConv(channels * 4, K=K, activation='relu', kernel_regularizer=l2(l2_reg), use_bias=False)([gcn1, A])
    gcn3 = ChebConv(channels * 8, K=K, activation='relu', kernel_regularizer=l2(l2_reg), use_bias=False)([gcn2, A])
    gcn4 = ChebConv(dataset.n_labels, K=K, activation='softmax', use_bias=True)([gcn3, A])

    # Build model
    model = Model(inputs=[X, A, I], outputs=gcn4)
    model.compile(optimizer=Adam(learning_rate), loss=CategoricalCrossentropy(reduction="sum"), weighted_metrics=['acc'])

# Fit data into model
model.summary()
model.fit(
    loader_tr.load(),
    batch_size=batch_size,
    steps_per_epoch=loader_tr.steps_per_epoch,
    epochs=epochs,
    shuffle=False,
    callbacks=[EarlyStopping(patience=patience, restore_best_weights=True)],
)

[danielegrattarola]

HI,

not really sure what's going on here. Does the crash happen if you train it on a single GPU?

I have worked with multi-GPUs settings in Spektral. in the past, but it's been a while so there may be bugs that need to be fixed.

[HRezaei]

Hi there,
It would be nice if you please add an example working with MirroredStrategy.
I tried a code like this but it doesn't work:

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  spektral_train_dataset = mirrored_strategy.experimental_distribute_dataset(spektral_train_dataset)

It raises an error like this:

AttributeError: 'TUDataset' object has no attribute '_variant_tensor'