Autograd yielding None in a correct computational graph

username: LDN
cooked: <p>Hi everyone! I’m trying to implement LISTA network (deepunfolded Iterative Soft Thresholding): in order to do so I implemented a custom layer and a network as follows:</p>
<p>Custom layer:</p>
<pre><code class="lang-auto">class LISTA_LAYER(nn.Module):
    def __init__(self, A):
        super().__init__()

        # Initialization of the learnable parameters

        self.W1 = nn.Linear(A.shape[0], A.shape[1], bias=False)
        self.W2 = nn.Linear(A.shape[1], A.shape[1], bias=False)

        self.beta = nn.Parameter(torch.ones(1, 1, 1), requires_grad=True)       
        self.mu = nn.Parameter(torch.ones(1, 1, 1), requires_grad=True)  
        
        # Apply Xavier uniform initialization
        nn.init.xavier_uniform_(self.W1.weight)
        nn.init.xavier_uniform_(self.W2.weight)

    def _shrink(self, x, beta):
        return beta * F.softshrink(x / beta, lambd=1)
    
    def forward(self, y, x):       
        if x is None:
            x = torch.zeros(y.shape[0], A.shape[1], requires_grad=True)

        return self._shrink(x - self.mu*(self.W1(y) + self.W2(x)), 
                            self.beta)
</code></pre>
<p>Custom network:</p>
<pre><code class="lang-auto">class LISTA_Net(nn.Module):
    def __init__(self, A, K = 5):
        super(LISTA_Net, self).__init__()

        # Number of layers &lt;-&gt; iterations
        self.K = K

        # Layers
        self.LISTA_layers = nn.ModuleList([LISTA_LAYER(A) for _ in range(self.K + 1)])

    def forward(self, y):     
        x = self.LISTA_layers[0].forward(y,None)

        for i in range(1, self.K + 1):
            x = self.LISTA_layers[i].forward(y,x)

        return x
</code></pre>
<p>The idea of LISTA is embodied in the following computational graph:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/3/4/34ee946f1b5950af9c69bd66a6fb96dca97ff391.png" alt="image" data-base62-sha1="7yfYY4QVlUv7HBxBNm5hUsCyt7r" width="573" height="166"></p>
<p>When I work out the main train loop I have None gradients everywhere, could you please help me figure out why? I mean, I checked that all the variables are seen as leaves of the computational graph looking at</p>
<p><code>list([p for p in model.parameters() if p.is_leaf])</code></p>
<p>at seing that all network parameters are in this list.</p>
<p>This is the main train loop once I instantiated the model:</p>
<pre><code class="lang-auto">optimizer = torch.optim.SGD(
        model.parameters(),
        lr=5e-05,
        momentum=0.9,
        weight_decay=0,
    )
    scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer, step_size=50, gamma=0.1
    )
    loss_train = np.zeros((num_epochs,))
    loss_test = np.zeros((num_epochs,))
    # Main loop
    for epoch in range(num_epochs):
        model.train()
        train_loss = 0
        for _, (b_x, b_s) in enumerate(train_loader):
            s_hat = model.forward(b_x)
            loss = F.mse_loss(s_hat, b_s, reduction="sum")
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            model.zero_grad()
            train_loss += loss.data.item()
        loss_train[epoch] = train_loss / len(train_loader.dataset)
        scheduler.step()
</code></pre>
<p>Thanks in advance <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <p>I cannot reproduce the issue using your model definition and this minimal code snippet:</p>
<pre data-code-wrap="python"><code class="lang-python">A = torch.randn(1, 1)
model = LISTA_Net(A)
x = torch.randn(1, 1)
out = model(x)
out.mean().backward()

for name, param in model.named_parameters():
    print(name, param.grad.abs().sum())

# LISTA_layers.0.beta tensor(0.3581)
# LISTA_layers.0.mu tensor(0.7568)
# LISTA_layers.0.W1.weight tensor(0.5244)
# LISTA_layers.0.W2.weight tensor(0.)
# LISTA_layers.1.beta tensor(0.2156)
# LISTA_layers.1.mu tensor(0.3776)
# LISTA_layers.1.W1.weight tensor(0.3157)
# LISTA_layers.1.W2.weight tensor(0.2401)
# LISTA_layers.2.beta tensor(0.0812)
# LISTA_layers.2.mu tensor(0.1312)
# LISTA_layers.2.W1.weight tensor(0.1190)
# LISTA_layers.2.W2.weight tensor(0.1515)
# LISTA_layers.3.beta tensor(0.0349)
# LISTA_layers.3.mu tensor(0.1621)
# LISTA_layers.3.W1.weight tensor(0.0512)
# LISTA_layers.3.W2.weight tensor(0.0867)
# LISTA_layers.4.beta tensor(0.1457)
# LISTA_layers.4.mu tensor(0.7042)
# LISTA_layers.4.W1.weight tensor(0.2134)
# LISTA_layers.4.W2.weight tensor(0.8921)
# LISTA_layers.5.beta tensor(1.)
# LISTA_layers.5.mu tensor(1.8967)
# LISTA_layers.5.W1.weight tensor(1.4645)
# LISTA_layers.5.W2.weight tensor(0.2895)
</code></pre>
<p>Could you post a minimal and executable code snippet reproducing the issue, please?</p>