How to combine NeighborLoader with muti-gpu?

[kkclear]

I am operating on a single graph, so I first split the graph into partitions using Neighborloader. Next, I want to know how to use muti-GPU. Here are my codes. How can I modify my code to meet multi-GPU？And Do I also need to use Neighborloader in my test_loader？ because when I evaluate model in test_loader (singer garph) , I wonder if it may cause out of memory. thanks a lot

dataset #my dataset like this
#Data(x=[871, 2000], edge_index=[2, 566903], y=[871], edge_weight=[566903], train_mask=[871], test_mask=[871])

from torch_geometric.loader import NeighborLoader
train_loader = NeighborLoader(dataset, num_neighbors=[25, 10], batch_size=16,
                        input_nodes=dataset.train_mask, shuffle=True)

from torch_geometric.data import DataLoader
test_loader = DataLoader([dataset], shuffle=True)  
# Do I also need to use **Neighborloader** in my test_loader? #single graph

for batch in train_loader:
   print(batch)
#Data(x=[800, 2000], edge_index=[2, 3410], y=[800], edge_weight=[3410], train_mask=[800], test_mask=[800], batch_size=16)
#Data(x=[861, 2000], edge_index=[2, 3590], y=[861], edge_weight=[3590], train_mask=[861], test_mask=[861], batch_size=16)
# ........
#Data(x=[854, 2000], edge_index=[2, 3620], y=[854], edge_weight=[3620], train_mask=[854], test_mask=[854], batch_size=16)
#Data(x=[802, 2000], edge_index=[2, 3410], y=[802], edge_weight=[3410], train_mask=[802], test_mask=[802], batch_size=16)

def train_model:
    for epoch in range(self.epochs):
        model.train()
        for data in self.train_loader:
            data = data.to(self.device)
            y_pred = model(data.x, data.edge_index, data.edge_weight)
            if self.use_mask:
                train_mask = data.train_mask.bool()
                y = data.y[train_mask]
                y_pred = y_pred[train_mask]
            else:
                y = data.y
            loss = loss_fn(y_pred, y)
            loss_train.append(loss.item())
            # backward
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            #saving gpu memory
            del loss
    return loss_train

def eval_model:
    model.eval()
    eval_fn = F1Score() if self.multi_label else Accuracy()
    for data in self.test_loader:
        data = data.to(self.device)
        out = model(data.x, data.edge_index)
        y = data.y
        if self.multi_label:
            y_pred = (out>0)
        else:
            y_pred = out.argmax(dim=1)
        if self.use_mask:
            test_mask = data.test_mask.bool()
            y = y[test_mask]
            y_pred = y_pred[test_mask]
        eval_fn(y_pred.int().view(-1).cpu(), y.int().view(-1).cpu())
    return eval_fn.compute().item()

[rusty1s]

The idea is that you would have a loader for every GPU you are training on. Here are two examples:

• https://github.com/pyg-team/pytorch_geometric/blob/master/examples/multi_gpu/distributed_sampling.py#L70
• https://github.com/pyg-team/pytorch_geometric/blob/master/examples/pytorch_lightning/graph_sage.py