PyTorch Geometric GIN-Conv layers parameters not updating

abc50111: <p>I made a composite model <code>MainModel</code> which consist of a <code>GinEncoder</code> and a <code>MainModel</code> which containing some <code>Linear</code> layers, and the <code>GinEncoder</code> made by the package <code>torch-geometric</code>, show as following codes :</p>
<pre><code class="lang-python">class GinEncoder(torch.nn.Module):
    def __init__(self):
        super(GinEncoder, self).__init__()
        self.gin_convs = torch.nn.ModuleList()
        self.gin_convs.append(GINConv(Sequential(Linear(1, 4),
                                                 BatchNorm1d(4), ReLU(),
                                                 Linear(4, 4), ReLU())))
        self.gin_convs.append(GINConv(Sequential(Linear(4, 4),
                                                 BatchNorm1d(4), ReLU(),
                                                 Linear(4, 4), ReLU())))


    def forward(self, x, edge_index, batch_node_id):
        # Node embeddings
        nodes_emb_layers = []
        for i in range(2):
            x = self.gin_convs[i](x, edge_index)
            nodes_emb_layers.append(x)

        # Graph-level readout
        nodes_emb_pools = [global_add_pool(nodes_emb, batch_node_id) for nodes_emb in nodes_emb_layers]

        # Concatenate and form the graph embeddings
        graph_embeds = torch.cat(nodes_emb_pools, dim=1)
        return graph_embeds


    def get_embeddings(self, x, edge_index, batch_node_id):
        with torch.no_grad():
            graph_embeds = self.forward(x, edge_index, batch_node_id).reshape(-1)

        return graph_embeds


class MainModel(torch.nn.Module):
    def __init__(self, graph_encoder:torch.nn.Module):
        super(MainModel, self).__init__()
        self.graph_encoder = graph_encoder
        self.lin1 = Linear(8, 4)
        self.lin2 = Linear(4, 8)


    def forward(self, x, edge_index, batch_node_id):
        graph_embeds = self.graph_encoder(x, edge_index, batch_node_id)
        out_lin1 = self.lin1(graph_embeds)
        pred = self.lin2(out_lin1)[-1]

        return pred

gin_encoder = GinEncoder().to("cuda")
model =  MainModel(gin_encoder).to("cuda")
</code></pre>
<p>I found that the weights of <code>GinEncoder</code> were not updated, while the weights of <code>Linear</code> layer in <code>MainModel</code> were updated.I observe this by following codes:</p>
<pre><code class="lang-python">gin_encoder = GinEncoder().to("cuda")
model =  MainModel(gin_encoder).to("cuda")
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())
epochs = 

for epoch_i in range(epochs):
    model.train()
    train_loss = 0

    for batch_i, data in enumerate(train_loader):
        data.to("cuda")
        x, x_edge_index, x_batch_node_id = data.x, data.edge_index, data.batch
        y, y_edge_index, y_batch_node_id = data.y[-1].x, data.y[-1].edge_index, torch.zeros(data.y[-1].x.shape[0], dtype=torch.int64).to("cuda")
        optimizer.zero_grad()
        graph_embeds_pred = model(x, x_edge_index, x_batch_node_id)
        y_graph_embeds = model.graph_encoder.get_embeddings(y, y_edge_index, y_batch_node_id)
        loss =  criterion(graph_embeds_pred, y_graph_embeds)
        train_loss += loss
        loss.backward()
        optimizer.step()
        if batch_i == 0:
            print(f"NO. {epoch_i} EPOCH")
            print(f"MainModel weights in epoch_{epoch_i}_batch0:{next(islice(model.parameters(), 15, 16))}", end="\n\n")
            print(f"GinEncoder weights in epoch_{epoch_i}_batch0:{next(model.graph_encoder.parameters())}")
            print("*"*80)
</code></pre>
<p>Outputs of codes:</p>
<pre><code class="lang-auto">NO. 0 EPOCH
MainModel weights in epoch_0_batch0:Parameter containing:
tensor([-0.1447, -0.3689, -0.2840, -0.3619, -0.2040,  0.2430,  0.4651,  0.3736],
       device='cuda:0', requires_grad=True)

GinEncoder weights in epoch_0_batch0:Parameter containing:
tensor([[-0.8312],
        [-0.5712],
        [-0.6963],
        [-0.1601]], device='cuda:0', requires_grad=True)
********************************************************************************
NO. 1 EPOCH
MainModel weights in epoch_1_batch0:Parameter containing:
tensor([-0.1842, -0.3333, -0.3170, -0.3247, -0.2424,  0.2627,  0.4272,  0.4119],
       device='cuda:0', requires_grad=True)

GinEncoder weights in epoch_1_batch0:Parameter containing:
tensor([[-0.8312],
        [-0.5712],
        [-0.6963],
        [-0.1601]], device='cuda:0', requires_grad=True)
********************************************************************************
NO. 2 EPOCH
MainModel weights in epoch_2_batch0:Parameter containing:
tensor([-0.2302, -0.3077, -0.3251, -0.2905, -0.2847,  0.2558,  0.3881,  0.4527],
       device='cuda:0', requires_grad=True)

GinEncoder weights in epoch_2_batch0:Parameter containing:
tensor([[-0.8312],
        [-0.5712],
        [-0.6963],
        [-0.1601]], device='cuda:0', requires_grad=True)
********************************************************************************
</code></pre>
<p>My question is how to make <code>loss.backward()</code> and <code>optimizer.step()</code> also pass to <code>GinEncoder</code>?</p>
<p>PS.</p>
<ul>
<li>I put the complete codes in here: <a href="https://gist.github.com/theabc50111/3ca708d0c1101d57b6172bd717302710" class="inline-onebox" rel="noopener nofollow ugc">a composite model composed of pytorch and torch-geometric · GitHub</a>
</li>
<li>I put the training data on Google Drive: <a href="https://drive.google.com/drive/folders/1_KMwCzf1diwS4gGNdSSxG7bnemqQkFxI?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">tmp - Google Drive</a>
</li>
</ul>

ptrblck: <p>Could you check the <code>.grad</code> attribute of all parameters of the <code>GinEncoder</code> before and after the first <code>.backward</code> call to see if these gradients are calculated but might be small?</p>

abc50111: <p>Thank you for your patience in reading my question.</p>
<p>I observe that the <strong>gradient of parameters Gin Model is always 0</strong>.</p>
<p>I tried to use following codes to observe the <code>.grad</code> attribute of <em>the first layer of parameters of the <code>GinEncoder</code></em>:</p>
<pre><code class="lang-auto">        x, x_edge_index, x_batch_node_id = data.x, data.edge_index, data.batch
        y, y_edge_index, y_batch_node_id = data.y[-1].x, data.y[-1].edge_index, torch.zeros(data.y[-1].x.shape[0], dtype=torch.int64).to("cuda")
        model_optimizer.zero_grad()
        graph_embeds_pred = model(x, x_edge_index, x_batch_node_id)
        y_graph_embeds = model.graph_encoder.get_embeddings(y, y_edge_index, y_batch_node_id)
        loss =  criterion(graph_embeds_pred, y_graph_embeds)
        train_loss += loss
        print(f"Before loss.backward(), MainModel weights.grad in epoch_{epoch_i}_batch{batch_i}:{next(islice(model.parameters(), 15, 16)).grad}", end="\n\n")
        print(f"Before loss.backward(), MainModel.graph_encoder weights.grad in epoch_{epoch_i}_batch{batch_i}:{next(model.graph_encoder.parameters()).grad}")
        loss.backward()
        print(f"After loss.backward(), MainModel weights.grad in epoch_{epoch_i}_batch{batch_i}:{next(islice(model.parameters(), 15, 16)).grad}", end="\n\n")
        print(f"After loss.backward(), MainModel.graph_encoder weights.grad in epoch_{epoch_i}_batch{batch_i}:{next(model.graph_encoder.parameters()).grad}")
        print("*"*80)
</code></pre>
<p>The result in epoch 0 &amp; <strong>batch 0</strong>:</p>
<pre><code class="lang-auto">Before loss.backward(), MainModel weights.grad in epoch_0_batch0:None

Before loss.backward(), MainModel.graph_encoder weights.grad in epoch_0_batch0:None
After loss.backward(), MainModel weights.grad in epoch_0_batch0:tensor([-0.0839,  0.0596, -0.1096,  0.0718,  0.1150,  0.0749,  0.0800,  0.0076],
       device='cuda:0')

After loss.backward(), MainModel.graph_encoder weights.grad in epoch_0_batch0:tensor([[0.],
        [0.],
        [0.],
        [0.]], device='cuda:0')
</code></pre>
<p>The result in epoch 0 &amp; <strong>batch 1</strong>:</p>
<pre><code class="lang-auto">Before loss.backward(), MainModel weights.grad in epoch_0_batch1:tensor([0., 0., 0., 0., 0., 0., 0., 0.], device='cuda:0')

Before loss.backward(), MainModel.graph_encoder weights.grad in epoch_0_batch1:tensor([[0.],
        [0.],
        [0.],
        [0.]], device='cuda:0')
After loss.backward(), MainModel weights.grad in epoch_0_batch1:tensor([-0.0640,  0.0315, -0.0785,  0.0666,  0.1209,  0.0641,  0.0495, -0.0090],
       device='cuda:0')

After loss.backward(), MainModel.graph_encoder weights.grad in epoch_0_batch1:tensor([[0.],
        [0.],
        [0.],
        [0.]], device='cuda:0')
</code></pre>

ptrblck: <p>OK, the results at least show that you are not detaching the operations from the computation graph since the <code>.grad</code> attributes are at least populated.<br>
I don’t know enough about the <code>GINConv</code> implementation to comment why the gradients might be zero, but <a class="mention" href="/u/rusty1s">@rusty1s</a> might know.</p>

abc50111: <p>Thank You for reply.</p>
<p>I have tried to merge the two class: <code>MainModel</code> and <code>GinEncoder()</code> to avoid the unchanging parameters of <code>gin_covs()</code>.</p>
<p>It works, it made the parameters of <code>gin_convs()</code> change, <strong>but this approach is still not answer my confusion</strong></p>
<p>Here is the update codes:</p>
<pre><code class="lang-auto">class MainModel2(torch.nn.Module):
    def __init__(self):
        super(MainModel2, self).__init__()
        self.gin_convs = torch.nn.ModuleList()
        self.gin_convs.append(GINConv(Sequential(Linear(1, 4), ReLU(),
                                                 Linear(4, 4), ReLU(),
                                                 BatchNorm1d(4))))
        self.gin_convs.append(GINConv(Sequential(Linear(4, 4), ReLU(),
                                                 Linear(4, 4), ReLU(),
                                                 BatchNorm1d(4))))
        self.lin1 = Linear(8, 4)
        self.lin2 = Linear(4, 8)


    def forward(self, x, edge_index, batch_node_id):
        # Node embeddings
        nodes_emb_layers = []
        for i in range(2):
            x = self.gin_convs[i](x, edge_index)
            nodes_emb_layers.append(x)

        # Graph-level readout
        nodes_emb_pools = [global_add_pool(nodes_emb, batch_node_id) for nodes_emb in nodes_emb_layers]

        # Concatenate and form the graph embeddings
        graph_embeds = torch.cat(nodes_emb_pools, dim=1)
        out_lin1 = self.lin1(graph_embeds)
        pred = self.lin2(out_lin1)[-1]

        return pred
</code></pre>

abc50111: <p>Update:<br>
The the weights are actually updated, I sums up all the <code>.grad</code> of every layers of <code>GinEncoder</code>.</p>
<p>I find that the <code>.grad</code> of the first layer of <code>GinEncoder</code> weights <strong>occasionally</strong> is zero, it has nothing to do with how you build the model from pytorch and torch_geometric.</p>
<p>So like <a class="mention" href="/u/ptrblck">@ptrblck</a> said, after check the <code>.grad</code> attribute of <strong>all</strong> parameters, the weights are actually updated, I just missed it.</p>