Multi GPU Training

[DuanhaoranCC]

Hello, I don't have enough GPU memory, but I have four GPU. I'm trying to change it to multi card training.

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Coauthor
from torch_geometric.nn import GCNConv

dataset = Coauthor(root='./', name="Physics")

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)

data = dataset[0]
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
data = data.to(device)
model = GCN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

I referred to it(https://github.com/pyg-team/pytorch_geometric/blob/master/examples/multi_gpu/data_parallel.py). But it is Graph Classfication, batch. I suggest adding a multi card training example for a single large graph. In addition, how can I modify it?

from torch_geometric.nn import DataParallel
model = DataParallel(model, device_ids=[0,1,2,3])

[rusty1s]

Please see here or here for an example on how to apply multi-GPU training on a single graph. DataParallel is only useful in case you want to split a dataset of graphs (not a single graph) onto multiple GPUs.

[DuanhaoranCC]

OK, thank you very much! In addition, I want to hear your opinion on why you based on pytorch_lightning instead of native pytoch when developing LightningNodeData. Finally, multi card training (single large graph) can only be based on sampling. Will this be lower than the performance of full graph?

[rusty1s]

LightningNodeData is just a convenient wrapper for people who want to make use of PyTorch Lightning (that properly handles data sharing across multiple GPUs). Native PyTorch works as well though. Due to sampling, there exist some variants in model performance compared to full-batch mode (the model might perform better or worse, but the difference should be minimal). If you want to replicate full-batch performance, you will have to sample all neighbors via num_neighbors = [-1] * num_layers (but keep in mind that this will increase memory consumption).