Training the adjacency matrix(weights) as a learnable parameter for graph classification

[BhuyanMonami]

Hello all,
I am trying to implement Graph CNN for a video data set where the adjacency matrix is based on the inverse Euclidean distance between the transformed features of the objects in a particular frame. Specifically, I have implemented the following:

def calculate_inverse_distances(features):
    in_channel = 256
    linear_layer = nn.Linear(in_channel, in_channel, bias=False)
    linear_layer.to(features.device)
    transformed_features = linear_layer(features)
    transformed_features = transformed_features.detach()
    num_features = transformed_features.shape[0]
    inverse_distances = torch.zeros((num_features, num_features))
    for i in range(num_features):
        for j in range(num_features):
            if i == j:
                continue
            # Calculate the Euclidean distance between feature vectors i and j
            distance = torch.dist(transformed_features[i], transformed_features[j],p=2)
            
            # Avoid division by zero by adding a small epsilon value
            epsilon = 1e-6
            inverse_distances[i, j] = 1.0 / (distance + epsilon)
            inverse_distances[j, i] = inverse_distances[i, j]  # Since distance(i, j) == distance(j, i)
    row_sums = inverse_distances.sum(dim=1, keepdim=True)
    normalized_inverse_distances = inverse_distances / row_sums
    normalized_inverse_distances = F.softmax(inverse_distances, dim=1)
    return inverse_distances

As you can see, I have applied a linear transformation to the feature vectors using nn.Linear and the weights of this layer are learnable parameters. I have used this adjacency matrix to create a data object. However, I am not sure if I can indeed treat a learnable parameter as a data object in PyG.
Specifically, I have done the following:

inverse_dists = calculate_inverse_distances(frame_window_features)
edge_index = inverse_dists.nonzero().t()
edge_weight = inverse_dists[edge_index[0], edge_index[1]]
edge_weight.requires_grad = True
data = Data(frame_window_features, edge_index=edge_index, edge_weight = edge_weight, y =labels)

My concerns regarding my approach are as follows:

1. Can the above approach indeed learn the adjacency matrix weights without additional manipulation or do I have to set the edge weights as Parameters separately? If not, that what would be an appropriate way to achieve my objective?
2. Since I intend to do batch-wise graph classification, with each graph having a variable number of nodes, can I treat learnable parameters (weights of the nn.Linear layer) as a data object?

Thanks in advance!

[ldv1]

That's interesting. I do not see why it should fail, but you will have I guess to remove

transformed_features = transformed_features.detach()

otherwise the parameters of your linear_layer will not be updated (with detach() you are setting the gradients to 0).
Furthermore, there is no need for

edge_weight.requires_grad = True

because this is not a parameter of your model.

[BhuyanMonami]

Hi. Thank you for your response. My concern was regarding using a learnable parameter(linear_layer in the calculate_inverse_distances in my case) as a graph data object. From your response, it seems like there shouldn't be any issue and the parameters of the linear_layer would still be updated if I use the data object in the downstream Gfraph CNN model.

Additionally, since I am trying to implement graph level classification, I have used dataloader to batch the input tensors like so:

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

And then for the model:

        self.sim_gc1 = GCNConv(in_channels, hidden_channels)
        self.sim_gc2 = GCNConv(hidden_channels, hidden_channels)
        self.sim_gc3 = GCNConv(hidden_channels, hidden_channels)
        self.lin = Linear(hidden_channels, out_channels)
        def forward(self, x,edge_index, edge_weight, batch):
        
            out = F.relu(self.sim_gc1(x, edge_index,edge_weight.sigmoid()))
            out = F.relu(self.sim_gc2(out, edge_index,edge_weight.sigmoid()))
            out = F.relu(self.sim_gc3(out, edge_index,edge_weight.sigmoid()))
            out = global_mean_pool(out, batch)
            out = F.dropout(out, p=0.5, training=self.training)
            out = self.lin(out)
            return out

(Here, I have simply put the code snippet of the forward method of my relevant model class). My specific concern is that when I use batching of graphs in the method shown, would it prevent the edge_weights (which are derived from the linear_layer of the features) from being updated?
Thanks!

[ldv1]

Now I understand the problem.

If you go for a batch size of 1, the following works:

import torch
from torch import nn
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
from torch_geometric.nn.conv import GCNConv
import torch.nn.functional as F
from torch_geometric.nn.pool import global_mean_pool

class Model(nn.Module):
    
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.linear_layer = nn.Linear(in_channels, in_channels, bias=False)
        self.sim_gc1 = GCNConv(in_channels, hidden_channels)
        self.sim_gc2 = GCNConv(hidden_channels, hidden_channels)
        self.sim_gc3 = GCNConv(hidden_channels, hidden_channels)
        self.lin = nn.Linear(hidden_channels, out_channels)
        
    def forward(self, x, edge_index, edge_weight, batch):
        out = F.relu(self.sim_gc1(x, edge_index,edge_weight.sigmoid()))
        out = F.relu(self.sim_gc2(out, edge_index,edge_weight.sigmoid()))
        out = F.relu(self.sim_gc3(out, edge_index,edge_weight.sigmoid()))
        out = global_mean_pool(out, batch)
        out = F.dropout(out, p=0.5, training=self.training)
        out = self.lin(out)
        return out
    
    def calculate_distances(self, features):
        # here just for the sake of illustration something simpler
        X = self.linear_layer(features)
        D = torch.cdist(X.unsqueeze(0), X.unsqueeze(0)).squeeze(0)
        return D
        
    def forward(self, x, edge_index, edge_weight, batch):
        out = F.relu(self.sim_gc1(x, edge_index,edge_weight.sigmoid()))
        out = F.relu(self.sim_gc2(out, edge_index,edge_weight.sigmoid()))
        out = F.relu(self.sim_gc3(out, edge_index,edge_weight.sigmoid()))
        out = global_mean_pool(out, batch)
        out = F.dropout(out, p=0.5, training=self.training)
        out = self.lin(out)
        return out

model = Model(2,10,3)

dataset = [ Data( torch.Tensor([ [1,1],[1, 1.1],[1,0.9] ]) ),
            Data( torch.Tensor([ [-1,-1],[1, 1.1],[1,0.9] ]) ),
            Data( torch.Tensor([ [1,1],[-1, -1.1],[1,0.9] ]) ) ]
loader = DataLoader(dataset, batch_size=1)

for step, data in enumerate(loader):
    x = data.x
    distances = model.calculate_distances(x)
    edge_index = distances.nonzero().t()
    edge_weight = distances[edge_index[0], edge_index[1]]
    batch = torch.zeros(len(x), device=x.device, dtype=torch.long)
    predictions = model(x, edge_index, edge_weight, batch)

The idea is to load one graph at a time, and for that graph create the adjacency matrix and edge weights based on the features.

Actually, you can put all the edge_index and edge_weight calculation in the forward method.

[BhuyanMonami]

Thanks again.
I have modified my code so that the linear layer is created and trained within the model (in the forward method). However, what I am not able to understand is why I should use a batch_size of 1 (i.e. load only one graph at a time) instead of using multiple graphs (min-batching) per batch?

[ldv1]

You can use more than a single graph per batch but then you will have to adapt calculate_inverse_distances. Right now, you are calculating the distances from all nodes to all others. That is ok for a single graph, but it breaks down for a batch because you do not want to compute the distance between two points from two different graphs.

[BhuyanMonami]

I see. That does make sense. However, I was under the assumption that using a readout layer as out = global_mean_pool(out, batch) takes care of it. As per PyG documentation (https://colab.research.google.com/drive/1I8a0DfQ3fI7Njc62__mVXUlcAleUclnb?usp=sharing#scrollTo=0gZ-l0npPIca). According to this documentation, DataLoader stacks the adjacency matrices in a diagonal fashion (creating a giant graph that holds multiple isolated subgraphs), and node and target features are simply concatenated in the node dimension. Am I missing something?
If so, do you think that you would be able to suggest how to adapt calculate_inverse_distances ?

[ldv1]

You can adapt your calculate_inverse_distances by using the information in the batch. See the documentation. batch is a vector of int. It has as many elements as the number of rows of the data tensor x and indicates for each element x[i] which graph it belongs to. With this information you can avoid computing the distance between two nodes that are not from the same graph. Some methods from PyG might be here of help, I do not know.