1
$\begingroup$

I am using Pytorch-Geometric library to implement a Graph Convolutional Layer(GCN) followed by few linear layers for a prediction task. But after training on graphs with np. of nodes being 10K and no. of edges being ~ 20K-40K for around 50 epochs in a batch-gradient descent manner, the gradient values hit close to zero values and show a terribly bad gradient flow during backpropagation. How can this be handled to ensure a proper gradient flow for weight training?

For reference, I've pasted the model code below:

import torch
from torch.nn.parameter import Parameter
import torch.optim as optim
import torch.nn.functional as F
import torch.nn as nn
from torch_geometric.nn import GCNConv

class Model(nn.Module):
    def __init__(self, nin=1, nhid1=128, nout=128, hid_l=64, out_l=1):
        super(Model, self).__init__()
       
        self.gc1 = GCNConv(in_channels= nin, out_channels= nhid1)
        self.gc2 = GCNConv(in_channels= nhid1, out_channels= nout)
        self.lay1 = nn.Linear(nout ,hid_l)
        self.l0 = nn.Linear(hid_l,hid_l)
        self.l1 = nn.Linear(hid_l,hid_l)
        self.lay2 = nn.Linear(hid_l ,out_l)
        self.active = nn.LeakyReLU(0.1)
       
        with torch.no_grad():
            self.gc1.weight = Parameter(nn.init.uniform_(torch.empty(nin,nhid1),a=0.0,b=1.0))
            self.gc1.bias = Parameter(nn.init.uniform_(torch.empty(nhid1),a=0.0,b=1.0))
            self.gc2.weight = Parameter(nn.init.uniform_(torch.empty(nhid1,nout),a=0.0,b=1.0))
            self.gc2.bias = Parameter(nn.init.uniform_(torch.empty(nout),a=0.0,b=1.0))
            self.lay1.weight = Parameter(nn.init.uniform_(torch.empty(hid_l, nout ),a=0.0,b=1.0))
            self.l0.weight = Parameter(nn.init.uniform_(torch.empty(hid_l, hid_l),a=0.0,b=1.0))
            self.l1.weight = Parameter(nn.init.uniform_(torch.empty(hid_l, hid_l),a=0.0,b=1.0))
            self.lay2.weight = Parameter(nn.init.uniform_(torch.empty(out_l,hid_l),a=0.0,b=1.0))
                       

    def forward(self, data):
        x, adj = data.x, data.edge_index
        x = self.active(self.gc1(x, adj))
        x = self.active(self.gc2(x, adj))
        x = self.active(self.lay1(x))
        x = self.active(self.l0(x))
        x = self.active(self.l1(x))
        x = self.active(self.lay2(x))
       
        return x

Below is the gradient flow across the layers:

enter image description here

$\endgroup$
2
  • $\begingroup$ Am I seeing correctly that your very last layer has zero grads, ie black line annotation? $\endgroup$ – hH1sG0n3 Jan 17 at 12:16
  • $\begingroup$ @hH1sG0n3 Yes. Also, the values of the gradients are extremely low. I removed the activation function from the last layer and there wasn't any difference in the grad flow. P.S.: I'm the OP. $\endgroup$ – Your IDE Jan 21 at 15:25

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.