Scroll left ⇢
# Load a dataset:
dataset = Planetoid(root, name="Cora")
# Create a mini-batch loader:
loader = NeighborLoader(dataset[0], num_neighbors=[25, 10])
# Create your GNN model:
class GNN(torch.nn.Module):
def __init__(self):
# Choose between different GNN building blocks:
self.conv1 = GCNConv(dataset.num_features, 16)
self.conv2 = GCNConv(16, dataset.num_classes)
def forward(self, x, edge_index):
x = self.conv1(x, edge_index).relu()
return self.conv2(x, edge_index)
# Train you GNN model:
for data in loader:
y_hat = model(data.x, data.edge_index)
loss = criterion(y_hat, data.y)
# Load a dataset:
dataset = TUDataset(root, name="PROTEINS")
# Create a mini-batch loader:
loader = DataLoader(dataset, batch_size=256)
# Create your GNN model:
class GNN(torch.nn.Module):
def __init__(self):
self.conv1 = GATConv(dataset.num_features, 16)
self.conv2 = GATConv(16, 16)
self.lin = Linear(16, dataset.num_classes)
def forward(self, x, edge_index, batch):
x = self.conv1(x, edge_index).relu()
x = self.conv2(x, edge_index).relu()
# Choose between different GNN building blocks:
x = global_mean_pool(x, batch)
return self.lin(x)
# Train you GNN model:
for data in loader:
y_hat = model(data.x, data.edge_index, data.batch)
loss = criterion(y_hat, data.y)
# Load a dataset:
dataset = Reddit(root)
# Create a mini-batch loader:
loader = LinkNeighborLoader(dataset[0], num_neighbors=[25, 10])
# Create your GNN model:
class GNN(torch.nn.Module):
def __init__(self):
# Choose between different GNN building blocks:
self.encoder = GraphSAGE(dataset.num_features, 16, num_layers=2)
Self.decoder = InnerProductDecoder()
def forward(self, x, edge_index, edge_label_index):
x = self.encoder(x, edge_index)
return self.decoder(x, edge_label_index)
# Train you GNN model:
for data in loader:
y_hat = model(data.x, data.edge_index, data.edge_label_index)
loss = criterion(y_hat, data.edge_label)