Training is slow on large graph

[garyhsu29]

Hi here!
My dataset has the following shape:
Data(x=[1838656, 62], edge_index=[2, 17746760], edge_attr=[17746760, 2], y=[1838656], train_mask=[1838656], val_mask=[1838656], test_mask=[1838656])

I want to train the node embedding for each node, so follow the example here: https://github.com/pyg-team/pytorch_geometric/blob/master/examples/graph_sage_unsup.py the only thing I change is the data and I also increase the batch_size of LinkNeighborLoader from 256 to 1024.

I move the data and model to my GPU (Tesla T4) for training, but it's still slow. it will take around 10 seconds for a batch, and I have 17331 batches in my training data.
0%| | 0/17331 [00:00<?, ?it/s].60.120.157)
0%| | 1/17331 [00:11<55:29:55, 11.53s/it])
0%| | 2/17331 [00:20<49:00:44, 10.18s/it])
0%| | 3/17331 [00:30<47:00:47, 9.77s/it])
0%| | 4/17331 [00:39<45:55:05, 9.54s/it])
0%| | 5/17331 [00:48<45:17:02, 9.41s/it])
0%| | 6/17331 [00:57<45:16:09, 9.41s/it])
0%| | 7/17331 [01:07<45:03:38, 9.36s/it])
0%| | 8/17331 [01:16<44:46:50, 9.31s/it])
0%| | 9/17331 [01:25<44:41:40, 9.29s/it])

I figure out that the training process only uses 1GB of my GPU memory, and GPU usage is around 1%.
Could you please provide any suggestions on this?

[garyhsu29]

I increased my batch_size in the LinkNeighborLoader to 1048576, and it's fast now.

[EdisonLeeeee]

Alternatively, you might take a look at Data Prefetching, which should increase your GPU usage.