Weights are not updaing? Does graph breaks when model passed through another class?

username: aswamy
cooked: <pre><code class="lang-auto">class NeuraNetworkClass(nn.Module):  # --&gt; This is learnable Neural Network
    def __init__(self,):
        super().__init__()
        self.fc1 = nn.Linear(.......)
        ..........
        ..........

    def forward(self, x)
        return self.fc1(x)


class NonNeuralNetworkClass(nn.Module):  # This is non learnable nn.Module class(no neural networks in this class)
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, a, b):
        return self.model(a) + self.model(b)

model = NeuraNetwork1()
net = NonNeuralNetworkClass(model)

optim = torch.optim.Adam(lr=1e-4, params=model.parameters())

for epoch in rannge(100):
    loss = loss_func(net(a=inp1, b=inp2), labels)
    optim.zero_grad()
    loss.backward()
    optim.step()
</code></pre>
<p><strong>Problem</strong>: With the above structure of coding, model parameters are not updating and gradients of model parameters are zero in every iteration.</p>
<p><strong>Question</strong>: I have passed the <code>model</code> as an argument to the init method of <code>NonNeuralNetworkClass</code>. So, does the forward method of <code>NonNeuralNetworkClass</code> gets the updated<br>
parameters of the <code>model</code> or do we need to pass the model every time in <code>forward</code> method of <code>NonNeuralNetworkClass</code> whenever the <code>model</code> parameter gets updated??</p>

username: ptrblck
cooked: <p>Your approach should work as seen here:</p>
<pre><code class="lang-python">class NeuraNetworkClass(nn.Module):
    def __init__(self,):
        super().__init__()
        self.fc1 = nn.Linear(1, 1)

    def forward(self, x):
        return self.fc1(x)


class NonNeuralNetworkClass(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, a, b):
        return self.model(a) + self.model(b)

model = NeuraNetworkClass()
net = NonNeuralNetworkClass(model)

optim = torch.optim.Adam(lr=1e-4, params=model.parameters())

loss_func = nn.MSELoss()
for epoch in range(10):
    inp1 = torch.randn(1, 1)
    inp2 = torch.randn(1, 1)
    labels = torch.randn(1, 1)
    loss = loss_func(net(a=inp1, b=inp2), labels)
    optim.zero_grad()
    loss.backward()
    for name, param in model.named_parameters():
        print(name, param.grad)
    optim.step()
</code></pre>
<p>so a zero gradient could indicate a vanishing gradient issue.</p>

username: aswamy
cooked: <p>Thanks for the response!</p>
<p>Weights of the model are never updated, I mean they are same as the initialised values. And the gradients are zero even from the 1st iteration. Can this still be a vanishing gradient problem?</p>

username: ptrblck
cooked: <p>Might be the case, but you should check, if the computation graph could be detached in some code, which wasn’t shown (it’s not detached in my example).<br>
To do so, create the model and make sure that ll <code>.grad</code> attributes are <code>None</code> via:</p>
<pre><code class="lang-python">for name, param in model.named_parameters():
    print(name, param.grad)
</code></pre>
<p>Skip the <code>optimizer.zero_grad()</code> call, create the model output and the loss, and call <code>loss.backward()</code>.<br>
Afterwards check the <code>.grad</code> attributes again.<br>
If they show zero values, the gradients were properly calculated in the <code>backward()</code> op but are zero e.g. due to the model architecture etc.<br>
If they are still <code>None</code>, the computation graph would be incorrect (assuming you expect all parameters to get a valid gradient).</p>

username: aswamy
cooked: <p>I tried this and got<code> params.grad</code> as None before <code>loss.backward()</code> and Zero after <code>backward()</code> step in the first iteration. If I continue for the next iterations, <code>params.grad</code> as zero both before and after as well, I think this is because in the first iteration grad is already zero and they are not updating.</p>
<p>Actually, to give a bit of an idea about architecture, I am trying to use the <code>model</code> to predict the parameters of the other network(let say <code>InferenceModel()</code>). The model architecture has 60 neural networks, each of which has 3 (inp, hid, out) layers. These 60 neural networks predicts(regress) 60 different parameters of <code>InferenceModel()</code>. This inference model takes input as images and outputs images(which means the loss is L2 between ground truth image and inferred image). So with this loss, I backpropagate the <code>model</code>'s parameters. So I backpropogate on the loss computed using the output of <code>InferenceModel()</code> to update <code>model</code> parameters.</p>
<p>Model architecture is shown below for 1st two networks out of 60. The same NN blocks will be repeated for remaininnig 58 NN as well.</p>
<pre><code class="lang-auto">(1): FCBlock(
      (net): Sequential(
        (0): FCLayer(
          (net): Sequential(
            (0): Linear(in_features=512, out_features=128, bias=True)
            (1): LeakyReLU(negative_slope=0.01, inplace=True)
          )
        )
        (1): FCLayer(
          (net): Sequential(
            (0): Linear(in_features=128, out_features=128, bias=True)
            (1): LeakyReLU(negative_slope=0.01, inplace=True)
          )
        )
        (2): FCLayer(
          (net): Sequential(
            (0): Linear(in_features=128, out_features=262144, bias=True)
            (1): Tanh()
          )
        )
      )
    )
    (2): FCBlock(
      (net): Sequential(
        (0): FCLayer(
          (net): Sequential(
            (0): Linear(in_features=512, out_features=128, bias=True)
            (1): LeakyReLU(negative_slope=0.01, inplace=True)
          )
        )
        (1): FCLayer(
          (net): Sequential(
            (0): Linear(in_features=128, out_features=128, bias=True)
            (1): LeakyReLU(negative_slope=0.01, inplace=True)
          )
        )
        (2): FCLayer(
          (net): Sequential(
            (0): Linear(in_features=128, out_features=512, bias=True)
            (1): Tanh()
          )
        )
      )
    )
</code></pre>
<p>FCBlock <a href="https://github.com/vsitzmann/metasdf/blob/56d0c6e411d23523bd638aa0ad37c0c2c9dbdb9d/meta_modules.py#L154" rel="noopener nofollow ugc">code</a><br>
FCLayer <a href="https://github.com/vsitzmann/metasdf/blob/56d0c6e411d23523bd638aa0ad37c0c2c9dbdb9d/meta_modules.py#L142" rel="noopener nofollow ugc">code</a></p>

username: aswamy
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I do not see anything wrong with my network. Gradients are still zero. Could you please help me debug this? I can share the complete code if it makes it easy for you to suggest something.</p>

username: ptrblck
cooked: <p>Sure, I can take a look at your code, if you could post a minimal and executable code snippet (using random input data) to reproduce this issue.</p>