比较图神经网络PyTorch Geometric 与 Deep Graph Library，帮助团队选出最适合的GNN库

virtualenv pyg_env –-python=python3 
source pyg_env/bin/activate 
pip install torch==1.10.0 
pip install torch-scatter torch-sparse torch-cluster \
   torch-spline-conv \
   torch-geometric -f https://data.pyg.org/whl/torch-1.10.0+cu102.html

import time 
import numpy as np 
import torch 
import torch.nn as nn 
import torch_geometric 
from torch_geometric.nn import GCNConv 
from torch_geometric.datasets import PPI 
from torch_geometric.loader import DataLoader
 
class GCNConvNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=6):
     super(GCNConvNet, self).__init__()
     self.gcn_0 = GCNConv(in_channels, 64)
     self.gcn_h1 = GCNConv(64, 64)
     self.gcn_h2 = GCNConv(64, 64)
     self.gcn_h3 = GCNConv(64, 64)
     self.gcn_h4 = GCNConv(64, 64)
     self.gcn_h5 = GCNConv(64, 64)
     self.gcn_h6 = GCNConv(64, 64)
     self.gcn_out = GCNConv(64, out_channels) 

def forward(self, batch):
     x, edge_index, batch_graph = batch.x, batch.edge_index, batch.batch
     x = torch.relu(self.gcn_0(x, edge_index))
     x = torch.relu(self.gcn_h1(x, edge_index))
     x = torch.relu(self.gcn_h2(x, edge_index))
     x = torch.relu(self.gcn_h3(x, edge_index))
     x = torch.relu(self.gcn_h4(x, edge_index))
     x = torch.relu(self.gcn_h5(x, edge_index))
     x = torch.relu(self.gcn_h6(x, edge_index))
     x = torch.dropout(x, p=0.25, train=self.training)
     x = self.gcn_out(x, edge_index)
     x = torch.sigmoid(x)
     return x

num_epochs = 10000 
lr = 1e-3 

dataset = PPI(root="./tmp") 
dataset = dataset.shuffle()
 
test_dataset = dataset[:2] 
train_dataset = dataset[2:] 

test_loader = DataLoader(test_dataset, batch_size=1) 
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True) 

for batch in train_loader:
     Break 

in_channels = batch.x.shape[1] 
out_channels = dataset.num_classes 

model = GCNConvNet(in_channels = in_channels,\
         out_channels = out_channels) 
loss_fn = torch.nn.BCELoss() 
optimizer = torch.optim.Adam(model.parameters(), lr=lr) 
my_device = "cuda" if torch.cuda.is_available() else "cpu" 
model = model.to(my_device)

losses = [] 
time_elapsed = [] 
epochs = [] 
t0 = time.time() 

for epoch in range(num_epochs): 

     total_loss = 0.0
     batch_count = 0
     for batch in train_loader: 
        optimizer.zero_grad() 
        pred = model(batch.to(my_device)) 
        loss = loss_fn(pred, batch.y.to(my_device)) 
        loss.backward() 
        optimizer.step() 
        total_loss += loss.detach() 
        batch_count += 1 

    mean_loss = total_loss / batch_count 
    losses.append(mean_loss) 
    epochs.append(epoch) 
    time_elapsed.append(time.time()-t0) 
    if epoch % 100 == 0: 
        print(f"loss at epoch {epoch} = {mean_loss}")

virtualenv dgl_env –python=python3 
source dgl_env/bin/activate 
pip install torch=1.10.0 

pip install dgl-cu102 -f https://data.dgl.ai/wheels/repo.html

import time 
import numpy as np 
import torch 
import torch.nn as nn 

import dgl 
from dgl.nn import GraphConv 
from dgl.data import PPIDataset 
from dgl.dataloading.pytorch import GraphDataLoader 
from dgl.data.utils import split_dataset 

class GraphConvNet(nn.Module): 

    def __init__(self, in_channels=3, out_channels=6): 
        super(GraphConvNet, self).__init__() 

        self.gcn_0 = GraphConv(in_channels, 64,\ 
        allow_zero_in_degree=True) 

        self.gcn_h1 = GraphConv(64, 64, allow_zero_in_degree=True) 
        self.gcn_h2 = GraphConv(64, 64, allow_zero_in_degree=True) 
        self.gcn_h3 = GraphConv(64, 64, allow_zero_in_degree=True) 
        self.gcn_h4 = GraphConv(64, 64, allow_zero_in_degree=True) 
        self.gcn_h5 = GraphConv(64, 64, allow_zero_in_degree=True) 
        self.gcn_h6 = GraphConv(64, 64, allow_zero_in_degree=True) 

        self.gcn_out = GraphConv(64, out_channels, \ 
        allow    _zero_in_degree=True)
 
    def forward(self, g, features): 

        x = torch.relu(self.gcn_0(g, features)) 

        x = torch.relu(self.gcn_h1(g, x)) 
        x = torch.relu(self.gcn_h2(g, x)) 
        x = torch.relu(self.gcn_h3(g, x)) 
        x = torch.relu(self.gcn_h4(g, x)) 
        x = torch.relu(self.gcn_h5(g, x)) 
        x = torch.relu(self.gcn_h6(g, x)) 

        x = torch.dropout(x, p=0.25, train=self.training) 
        x = self.gcn_out(g, x) 

        x = torch.sigmoid(x) 

        return x

if __name__ == "__main__": 
    num_epochs = 10000 
    lr = 1e-3 
    my_seed = 42 

    dataset = PPIDataset() 

    # randomly create dataset split 
    test_dataset, train_dataset = split_dataset(dataset, [0.1, 0.9],\ 
    shuffle=True) 
    train_loader = GraphDataLoader(train_dataset, batch_size=8) 
    test_loader = GraphDataLoader(test_dataset, batch_size=1) 

    for batch in train_loader: 
    break 

    in_channels = batch.ndata["feat"].shape[1] 
    out_channels = dataset.num_labels 
    model = GraphConvNet(in_channels, out_channels) 

    loss_fn = torch.nn.BCELoss() 
    optimizer = torch.optim.Adam(model.parameters(), lr=lr) 
    my_device = "cuda" if torch.cuda.is_available() else "cpu" 

    model = model.to(my_device) 

    losses = [] 
    time_elapsed = [] 
    epochs = [] 
    t0 = time.time() 

    for epoch in range(num_epochs): 

        total_loss = 0.0 
        batch_count = 0 

        for batch in train_loader: 
            optimizer.zero_grad() 
            batch = batch.to(my_device) 
            pred = model(batch, batch.ndata["feat"].to(my_device)) 
            loss = loss_fn(pred, batch.ndata["label"].to(my_device)) 
            loss.backward() 
            optimizer.step() 

            total_loss += loss.detach() 
            batch_count += 1 

        mean_loss = total_loss / batch_count 
        losses.append(mean_loss) 
        epochs.append(epoch) 
        time_elapsed.append(time.time() - t0) 

        if epoch % 100 == 0: 
            print(f"loss at epoch {epoch} = {mean_loss}") 

    # get test accuracy score 
    num_correct = 0. 
    num_total = 0. 
    model.eval() 

    for batch in test_loader: 

        batch = batch.to(my_device) 
        pred = model(batch, batch.ndata["feat"]) 
        num_correct += (pred.round() == \ 
            batch.ndata["label"].to(my_device)).sum() 
        num_total += pred.shape[0] * pred.shape[1] 

        np.save("dgl.npy", \ 
            {"epochs": epochs, \ 
            "losses": losses, \ 
            "time_elapsed": time_elapsed}) 

    print(f"test accuracy = {num_correct / num_total}")