problem with mini batching the adjacency matrix created from the nodes and edge_indices

[akpas001]

Hi,

I am trying to implement the concept of VGAE (variational graph autoencoder ) on a custom dataset. so, for that i have to create an adjacency matrix. The way i created my adjacency matrix is as follows.

I have 9 nodes per graph, and dynamic edge_indices. since I have 9 nodes. I have created a 9*9 matrix and checked for the active edges present.

dataset = []
adj = torch.zeros(9,9)
edges = edge_index[:,:int(active_edges)] # this gives the number of active edges present in that certain graph.
for i in edges.t():
     adj[i[0]i[1]] = 1 # replacing the zeros with ones whereever an active edge is present.
     dataset.append(Data(x =x, edge_index = edge_index[:,:int(active_edges)], adj _= adj)) # appending this inro a dataloader.

train_loader = DataLoader(train_dataset, batch_size = 32)

for data in train_loader:
    x, edge_index, adj = data.x, data.edge_index, data.adj
    print(x.shape) = [288,30] # because (batch_size*num_nodes, num_features)
    print(edge_index.shape) = [2, 441] # because (coo format) and  441 can vary for different because it is dynamic.
    print(adj.shape) = [288,9] # because again (batch_size*num_nodes, 9)

so, now the problem is during training. when i predict my adjacency matrix, using VGAE my latent feature output shape is (288, 20)
and i pass it through Innerproduct decoder to calculate adjacency and it returns me matrix of shape of (288,288) bacause torch.matmul(lat_fea, lat_fea.t()).

Now i cannot backpropogate these two calculate my loss. I dont understand if I am missing something here. Please help.

[Padarn]

Sorry I couldn't follow your example, what do you mean

Now i cannot backpropogate these two calculate my loss. I dont understand if I am missing something here. Please help.

Could you provide a working code block with your issue?

[akpas001]

sure. I did upload my code in the comment below. please feel free to go through it.

[rusty1s]

If you want to reconstruct edges batch-wise, you will need to first convert your node feature matrix of shape (288, 20) to (batch_size, num_nodes, 20) = (32, 9, 20). Then, you can re-construct edges for each adjacency matrix:

recon = x @ x.transpose(1, 2)

[akpas001]

Thankyou for taking out your time for helping me. So, basically the main objective of the VGAE model, I am implementing is to predict the adjacency matrix. so, to have a target_adjacency matrix, I have created an adjacency matrix using nodes and edge_indices. Now, when I am predicting the adjacency matrix. There is a size mismatch happening between the predicted adjacency (288,288) and target adjacency (288,9). (288 is because batch_size = 32 and num_nodes = 9). The reason for this is mini batching. So, my question is during the creation of adjacency matrix using nodes and edge_indices am I missing anything? I have uploaded my code and data for your reference in the below comment. Please feel free to go through it and help me.
Thank you!

[Padarn]

I think you just want to mask your target adjacency to only include the nodes in that minibatch?

[akpas001]

masking my target adjacency matrix doesnot convert it into the shape (288,288). Does it?

[Padarn]

Sorry was replying from my phone and missed some context.

I think what @rusty1s suggests is the best option (assuming I understand correctly): Your target_adjacency is a stack of 32, 9 x 9 adjacency graphs. So you can do:

target_adjacency = target_adjacency.reshape(32, 9, 9)

which will match the output share of what @rusty1s suggests

recon = x @ x.transpose(1, 2)

If you really want to turn your target_adjacency into the (288,288) shape, you can do something like

target_adjacency = torch.block_diag(*[target_adjacency[i*9:(i+1)*9, :] for i in range(32)])

which will build a block diagonal adjacency matrix, but this will require a lot more memory

[akpas001]

Thanks for the inputs. The first concept worked.

[akpas001]

states.zip
adjacency_pred_vgae .txt

Please feel free to go through the code and use the data provided, if it helps you to understand the problem batter. Thanks

[akpas001]

I just found a paper for Directed Graph Autoencoders where they are implementing Asymmetric Inner Prooduct Decoder for the decoder.
I am still going through it, it looks a bit complicated for me. Will get back to you if I can find some solution there. Else again will be back with my doubts. :D. I really appriciate your help!
https://www.aaai.org/AAAI22Papers/AAAI-5979.KolliasG.pdf

[akpas001]

According to this paper they are outputting two latent vectors from the encoder one for source and the other for target and calculating crossproduct between z_s, z_t.(t). How is that possible, i am having a problem when trying to do that with my edge_index. My model looks something of this sort . It's returning me an error on edge_index saying

RuntimeError: index 4 is out of bounds for dimension 0 with size 4.

please Help!

import numpy as np
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.loader import DataLoader
from torch_geometric.data import Data
from torch_geometric.nn import Linear, Sequential, GINConv, VGAE, GCNConv
from torch_geometric.nn.models import InnerProductDecoder
from dataloader import target_adj, states
import matplotlib.pyplot as plt
import time

# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')
batch_size = 32
t = time.time()
#%%
class VGAE(nn.Module):
    def __init__(self):
        super(VGAE,self).__init__()
        
        dim1 = 30
        dim2 = 45
        dim3 = 60
        dim4 = 40
        dim5 = 20
        dim6 = 15
        dim7 = 9
        
        self.mlp1 = nn.Sequential(Linear(dim1, dim2, weight_initializer=('glorot')),\
                                  nn.ReLU(), Linear(dim2, dim3))
                                 
        self.mlp2 = nn.Sequential(Linear(dim3, dim4),nn.Tanh(),Linear(dim4, dim5))
        
        self.mlp3 = nn.Sequential(Linear(dim1, dim2, weight_initializer=('glorot')),\
                                  nn.ReLU(), Linear(dim2, dim3))
                                 
        self.mlp4 = nn.Sequential(Linear(dim3, dim4),nn.ReLU(),Linear(dim4, dim5))
            
        self.mlp5 = nn.Sequential(Linear(dim5, dim4),nn.ReLU(),Linear(dim4, dim2))
            
        self.mlp6 = nn.Sequential(Linear(dim2, dim1),nn.ReLU(),Linear(dim1, dim5))
        
        self.mlp7 = nn.Sequential(Linear(dim5, dim6),nn.ReLU(),Linear(dim6, dim7))
        
        self.conv1 = GINConv(self.mlp1)
        self.conv2 = GINConv(self.mlp2)
        self.conv3 = GINConv(self.mlp3)
        self.conv4 = GINConv(self.mlp4)
        self.conv5 = GINConv(self.mlp5)
        self.conv6 = GINConv(self.mlp6)
        self.conv7 = GINConv(self.mlp7)
    def encode(self, data):
        
        x = data[:270].reshape(9,30).float()   
        edge_index = data[391:431].reshape(2,-1).long()
        act_edges = int(data[431])
        edge_index = edge_index[:,:int(act_edges)]
        
        src_x = x[edge_index[0]]
        tar_x = x[edge_index[1]]
        x = x.float()
        mu_src = self.conv1(src_x,edge_index) 
        log_var_src = nn.ReLU()(self.conv2(mu_src, edge_index))
        mu_tar = self.conv3(x,edge_index)
        log_var_tar = nn.ReLU()(self.conv4(log_var_src,edge_index))
        std_src = torch.exp(0.5*log_var_src) # reparameterization
        std_tar = torch.exp(0.5*log_var_tar)
        epsilon = torch.randn_like(std_src)
        z_src = mu_src + epsilon*std_src
        z_tar = mu_tar + epsilon*std_tar
        return z_src, z_tar, mu_src, log_var_src, mu_tar, log_var_tar
    
    def dot_prod_decode(self,z_src, z_tar):
        
        # edge_index = data[391:431].reshape(2,-1).long()
        # act_edges = int(data[431])
        # edge_index = edge_index[:,:int(act_edges)]
        
        # adj_pred = self.conv5(z, edge_index)
        # adj_pred = self.conv6(adj_pred, edge_index)
        # adj_pred = self.conv7(adj_pred, edge_index)
        adj_pred = torch.matmul(z_src,z_tar.t())

        return adj_pred
        
    def total_loss(self, adj_pred, adj, mu, log_var ):
        ns_loss = nn.BCEWithLogitsLoss()
        kld_loss = -0.5 * torch.mean(torch.sum(1+2*log_var - mu**2 -\
                                                log_var.exp()**2, dim = 1))
        recons_loss = ns_loss(adj_pred, adj)
        # re_loss = reg_loss(mu, log_var)
        total_loss = recons_loss + kld_loss  
        
        return total_loss, recons_loss, kld_loss 

#%% training. Not being able to shuffle data, gradients might be memorized. 
model = VGAE().to(device)
model_optim = optim.Adam(model.parameters(), lr = 8e-04)
model.train()
torch.autograd.set_detect_anomaly(True)
epochs = 250
tot_loss = []
tot_r_loss = []
tot_mse_loss = []
for ep in range(epochs):
    print('\rEpoch Number ', ep, end = '')
    ep_loss = 0
    ep_recon_loss = 0
    ep_reg_loss = 0
    count = 0
    for data, adj in zip(states, target_adj):
        count+=1
        # print('\rietration number is', count, end = '')
        z_src, z_tar,mu_src, log_var_src, mu_tar, log_var_tar = model.encode(data)
        adj_pred = model.dot_prod_decode(z_src,z_tar)
        loss, recons_loss, re_loss = model.total_loss(adj_pred, adj, mu_src,\
                                                      log_var_src, mu_tar, log_var_tar)
        # print('recons_loss', recons_loss.item())
        ep_loss += loss.item()
        ep_recon_loss += recons_loss.item()
        ep_reg_loss += re_loss.item()
    # print('\nno.of ietrations in this epoch are ', count, end = '')
    model_optim.zero_grad()
    loss.backward()
    # for name, params in model.named_parameters():
    #     print('gradients', params.grad.mean())
    #     print('weights', params.data)
    model_optim.step()

    tot_loss.append(ep_loss)
    tot_r_loss.append(ep_recon_loss)
    tot_mse_loss.append(ep_reg_loss)
    print("Epoch:", '%04d' % (ep + 1), "train_loss=", "{:.5f}".format(ep_loss),\
            "recon_loss=", "{:.5f}".format(ep_recon_loss),"time=", "{:.5f}".format(time.time() - t))
plt.title('tot_loss')
plt.plot(tot_loss)
plt.show()

plt.title('recons_loss')
plt.plot(tot_r_loss)
plt.show()

plt.title('kld loss')
plt.plot(tot_mse_loss)
plt.show()

[rusty1s]

Encode would still operate on the same node feature matrix, but will use different GNN layers for learning two representations which you can then lift to source and destination node representations.

[akpas001]

Thank you! That worked! But,as far as I can think we will not be able to apply the concept of variational graph autoencoders to this. Because, it is creating the problem of calculating KL Divergence loss for two latent vectors z1 and z2. Is there any way around it?

[rusty1s]

Why isn't that possible? The KL divergence is always only computed on individual vectors. Let me know if I am missing something.