How to correctly pass in adjacency matrix varaibles to graphSAGE?

[amytee252]

Hi,

I am having some trouble trying to pass in the size of the adjacency matrix to graphSage. This is the code as is:

data = PPI(root = "/GNN_Protein/ppi")[0]
shuf = np.random.choice(np.arange(data.x.shape[0]), data.x.shape[0], replace=False)
n_train = int(0.8 * 0.7 * shuf.shape[0])
n_test = int(0.8 * 0.3 * shuf.shape[0])
train_mask = torch.tensor([i in set(shuf[:n_train].tolist()) for i in np.arange(data.x.shape[0])], dtype=torch.bool)
test_mask = torch.tensor([i in set(shuf[n_train:(n_train+n_test)].tolist()) for i in np.arange(data.x.shape[0])], dtype=torch.bool)
holdout_mask = torch.tensor([i in set(shuf[(n_train+n_test):].tolist()) for i in np.arange(data.x.shape[0])], dtype=torch.bool)
print("Training on", sum(train_mask), "of",  shuf.shape[0])
print("Testing on", sum(test_mask), "of", shuf.shape[0])
print("Reserving", sum(holdout_mask), "of", shuf.shape[0])


reducer = umap.UMAP()
embedding = reducer.fit_transform(data.x.numpy()[:,0:5])
embedding.shape
snsumap = sns.scatterplot(embedding[:, 0],
	    embedding[:, 1])
snsumap.get_figure().savefig("plot.png")



# Assign labels based on edges per node.
cnts = data.edge_index[0].unique(return_counts=True)
node_counts = torch.zeros(data.x.shape[0], dtype=torch.long)
node_counts = node_counts.scatter_(0, cnts[0], cnts[1])
node_counts = node_counts.unsqueeze(-1)

# Create labels based on edge count, and otherwise random covariates.
x_random = torch.rand((node_counts.shape[0], 5))
y_simple = (node_counts > 15).float()
x_simple = torch.cat((x_random.float(), node_counts.float()), dim=1)

# Create dataset.
simple_data = torch_geometric.data.Data(edge_index=data.edge_index.clone(),
                                        x=x_simple,
                                        y=y_simple,
                                        train_mask=train_mask,
                                        test_mask=test_mask)
print(simple_data)

train_loader = NeighborSampler(simple_data.edge_index,
                                                    node_idx=torch.arange(simple_data.train_mask.shape[0])[simple_data.train_mask].long(),
                                                    sizes=[1],
                                                    batch_size=1,
                                                    shuffle=False,
                                                    num_workers=0)

# iterate over data.
for epoch in range(1):
    seen = 0
    for batch_size, n_id, adjs in train_loader:
        seen += batch_size
    print('batch_size ', seen, flush=True)

class SAGE(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers=3):
        super(SAGE, self).__init__()

        self.num_layers = num_layers

        self.convs = torch.nn.ModuleList()
        self.convs.append(SAGEConv(in_channels, hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(SAGEConv(hidden_channels, hidden_channels))
        self.convs.append(SAGEConv(hidden_channels, out_channels))

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()

    def forward(self, x, adjs):
            for i, (edge_index, _, size) in enumerate(adjs):
            xs = []
            print(xs)
            x_target = x[:size[1]]  # Target nodes are always placed first.
            print(x_target)
            x = self.convs[i]((x, x_target), edge_index)
            print(x)
            if i != self.num_layers - 1:
                x = F.relu(x)
                x = F.dropout(x, p=0.5, training=self.training)
            xs.append(x)
            if i == 0: 
                x_all = torch.cat(xs, dim=0)
                layer_1_embeddings = x_all
            elif i == 1:
                x_all = torch.cat(xs, dim=0)
                layer_2_embeddings = x_all
            elif i == 2:
                x_all = torch.cat(xs, dim=0)
                layer_3_embeddings = x_all    
        #return x.log_softmax(dim=-1)
        return layer_1_embeddings, layer_2_embeddings, layer_3_embeddings

    def inference(self, x_all):
        pbar = tqdm(total=x_all.size(0) * self.num_layers)
        pbar.set_description('Evaluating')
        total_edges = 0
        for i in range(self.num_layers):
            xs = []
            for batch_size, n_id, adj in subgraph_loader:
                edge_index, _, size = adj.to(device)
                total_edges += edge_index.size(1)
                x = x_all[n_id].to(device)
                x_target = x[:size[1]]
                x = self.convs[i]((x, x_target), edge_index)
                if i != self.num_layers - 1:
                    x = F.relu(x)
                xs.append(x)

                pbar.update(batch_size)

            if i == 0: 
                x_all = torch.cat(xs, dim=0)
                layer_1_embeddings = x_all
            elif i == 1:
                x_all = torch.cat(xs, dim=0)
                layer_2_embeddings = x_all
            elif i == 2:
                x_all = torch.cat(xs, dim=0)
                layer_3_embeddings = x_all
                
        pbar.close()

        return layer_1_embeddings, layer_2_embeddings, layer_3_embeddings



device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SAGE(50, 256, 10, num_layers=3)
model = model.to(device)

x = data.x.to(device)
y = data.y.squeeze().to(device)

def train(epoch):
    model.train()

    #pbar = tqdm(total=train_idx.size(0))
    #pbar.set_description(f'Epoch {epoch:02d}')

    total_loss = total_correct = 0
    for batch_size, n_id, adjs in train_loader:
        # `adjs` holds a list of `(edge_index, e_id, size)` tuples.
        #adjs = [adj.to(device) for adj in adjs]
        optimizer.zero_grad()    
        l1_emb, l2_emb, l3_emb = model(x[n_id], adjs)
        #print("Layer 1 embeddings", l1_emb.shape)
        #print("Layer 2 embeddings", l1_emb.shape)
        out = l3_emb.log_softmax(dim=-1)
        loss = F.nll_loss(out, y[n_id[:batch_size]])
        loss.backward()
        optimizer.step()

        total_loss += float(loss)
        total_correct += int(out.argmax(dim=-1).eq(y[n_id[:batch_size]]).sum())
        #pbar.update(batch_size)

    #pbar.close()

    loss = total_loss / len(train_loader)
    approx_acc = total_correct / train_idx.size(0)

    return loss, approx_acc


optimizer = torch.optim.Adam(model.parameters(), lr=0.003)

for epoch in range(1, 21):
    loss, acc = train(epoch)
    print(f'Epoch {epoch:02d}, Loss: {loss:.4f}, Approx. Train: {acc:.4f}')

I get this error:

Traceback (most recent call last):
File "main_final.py", line 260, in
loss, acc = train(epoch)
File "main_final.py", line 237, in train
l1_emb, l2_emb, l3_emb = model(x[n_id], adjs)
File "/home/teea1/myenv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "main_final.py", line 153, in forward
for i, (edge_index, _, size) in enumerate(adjs):
ValueError: not enough values to unpack (expected 3, got 2)

I honestly have no idea why it is not able to unpack everything in the adjacency matrix, I think it is having a problem with size.

[rusty1s]

I am sorry for the confusion. If the NeighborSampler will only sample from a single hop, then adjs will not be a list but will instead already return the adjacency matrix. As such, you can either drop the for-loop inside forward, or write:

        if not isinstance(adjs, list):
            adjs = [adjs]

at the start of it.

[amytee252]

Thank you so much for your response, it solved the problem