MNIST example.

[cnhung]

1. Why GPU is running 7.5 times faster than 8TPUs.
2. How about add MNIST_SIZE to _mnist_grid_graph(k, MNIST_SIZE) such that it can be imported with variable MNIST_SIZE.

[danielegrattarola]

Hi,

Why GPU is running 7.5 times faster than 8TPUs.

Depending on your setup, it could be for a number of reasons. I'd need more details to help you.

How about add MNIST_SIZE to _mnist_grid_graph(k, MNIST_SIZE) such that it can be imported with variable MNIST_SIZE.

MNIST_SIZE is meant to be a constant, since the MNIST dataset has 28 x 28 images. If you're rescaling the images, you can write your own dataset to account for this change, by extending the spektral.data.Dataset class. It should be quick to do.

Cheers

[cnhung]

https://www.kaggle.com/cnhung/mnist-p-graphnn-spektral

def get_strategy():
gpu = ""
try:
tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', tpu.master())
except ValueError:
tpu = None
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
gpu = tf.config.list_physical_devices("GPU")
if len(gpu) == 1:
print('Running on GPU ', gpu)
if tpu:
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu)
elif len(gpu) == 1:
strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
tf.config.optimizer.set_experimental_options({"auto_mixed_precision":True})
else:
strategy = tf.distribute.get_strategy()
print("REPLICAS: ", strategy.num_replicas_in_sync)
return strategy

strategy = get_strategy()

. . .

with strategy.scope():
if LANGUAGE==0:
model = Net()
else:
model = Net_Russian() #need further hyperparams tuning

optimizer = Adam(lr=0.001)
loss_fn = SparseCategoricalCrossentropy()
acc_fn = SparseCategoricalAccuracy()

[danielegrattarola]

So the fact that the computation is running ~8 times slower using 8 TPUs makes me wonder whether you're not repeating the same computation 8 times?
I don't think this issue is related to Spektral, anyway. You can try the following to make sure: remove any layer imported from Spektral and re-run the computation. If it's still slower, then you have to debug the distribution strategy.

Cheers