How to use multiple GPU when using pyG?

[hantongyou]

I am a beginer using pyG, and I have created my own database, which has 110275 nodes and each nodes got a 1536 dims feature, and the tagged nodes is about 6906. I tried to allocate the data and my GNN model(GAT) to the GPU but I failed. It tells me that CUDA out of memory, but I have no idea about how to use multiple GPU when running pyG code, could you please give me some advice? Thanks!

I referenced the sample here

class GAT(nn.Module):
    def __init__(self, hidden_channels) -> None:
        super().__init__()
        self.conv1 = GATConv(data.num_features,hidden_channels)
        self.conv2 = GATConv(hidden_channels,9)
        self.activation = nn.ReLU()
    def forward(self,data):
        x = data.x
        # print(data.x.shape)
        edge_index = data.edge_index
        x = self.conv1(x,edge_index)
        x = self.activation(x)
        x = self.conv2(x,edge_index)
        return x
model = GAT(hidden_channels=32)
model.cuda()
model.to(device)
if torch.cuda.device_count() > 1:
     print("Let's use", torch.cuda.device_count(), "GPUs!")
     model = DataParallel(model
                          , device_ids=[0,1,2,3]
                          )

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scheduler = lr_scheduler.StepLR(optimizer, 100, 0.8)
criterion = ContrastiveLoss()
def train():
    for batch_idx,batch in enumerate(train_loader):
        model.train()
        print(batch)
        optimizer.zero_grad() 
        out = model([data])  
        loss = criterion(data.x[data.train_mask], data.y[data.train_mask],out[data.train_mask])  
        loss.backward()
        optimizer.step() 
    return loss

def test():
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1) 
    test_correct = pred[data.train_mask] == data.y[data.train_mask]  
    test_acc = int(test_correct.sum()) / int(data.train_mask.sum()) 
    return test_acc

accs = []
for epoch in range(1, 1025):
    loss = train()
    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}',end=" ")
    acc = test()
    print("acc:",acc)
    accs.append(acc)
    scheduler.step()

[LukeLIN-web]

Currently, we have pytorch_geometric/examples/multi_gpu/ examples.

[hantongyou]

Thanks, I will try.

[LukeLIN-web]

Thanks, I will try.

If you have further problems, feel free to talk about them.