Specify retain_graph=True when calling backward the first time : differential-machine-learning

username: pascal_notsawo
cooked: <p>Hello.<br>
I have a multilayer perceptron made of several linear layers and an activation function to be specified when creating the model. Here’s how I implemented this (this is more or less the reimplementation of <a href="https://github.com/differential-machine-learning/notebooks/blob/master/DifferentialML.ipynb" rel="nofollow noopener">this</a> in pytorch):</p>
<pre><code class="lang-auto">class Linear(nn.Module):
    def __init__(self, in_features, out_features, bias = True, omega_0 = 1, activation_function = None, deriv_activation_function = None):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features, bias=bias)
        self.activation_function = activation_function if activation_function else lambda x : x
        self.deriv_activation_function = deriv_activation_function if deriv_activation_function else lambda x : x
        self.omega_0 = omega_0
  
    def forward(self, x):
        z = self.omega_0 * self.linear(x)
        y = self.activation_function(z)
        return y, z

class MLP(nn.Module):
    def __init__(self, in_features, hidden_features, hidden_layers, out_features, activation_function = None, deriv_activation_function = None):
        super().__init__()
        
        self.net = []

        # input layer
        self.net.append(Linear(in_features, hidden_features, activation_function, deriv_activation_function = deriv_activation_function))
    
        # hidden layer(s)
        for i in range(hidden_layers):
            self.net.append(Linear(hidden_features, hidden_features, activation_function, deriv_activation_function = deriv_activation_function))

        # output layer
        self.net.append(Linear(hidden_features, out_features,))

        self.net = nn.Sequential(*self.net)
        self.zs = []

    def forward(self, x):
        x_grad = x.clone().detach().requires_grad_(True) # allows to take derivative w.r.t. input
        for linear_layer in self.net :
             x, z = linear_layer(x)
             self.zs.append(z)
        y = x
        return y, x_grad 
    
    def backprop(self, y):
        zbar = torch.ones_like(y)
        for l in range(len(self.net) - 1, 0, -1):
            linear_layer = self.net[l] 
            zbar =  torch.matmul(zbar, linear_layer.linear.weight) * linear_layer.deriv_activation_function(self.zs[l-1])
        
        linear_layer = self.net[0]
        zbar = torch.matmul(zbar, linear_layer.linear.weight)
        
        xbar = zbar

        # dz[L] / dx
        return xbar
</code></pre>
<p>The backprop method just takes the differential of the output (y) with respect to itself (i.e. 1) and passes it from the last to the first layer of the neural network (the author called it <a href="https://arxiv.org/abs/2005.02347" rel="nofollow noopener">twin_net</a>).<br>
During training, as shown below, I get this error.</p>
<pre><code class="lang-auto">RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling backward the first time.
</code></pre>
<p>How can I correct it? I was told <a href="https://stackoverflow.com/questions/48274929/pytorch-runtimeerror-trying-to-backward-through-the-graph-a-second-time-but" rel="nofollow noopener">here</a> to play with .detach(), but I’m not sure which variable to play with.</p>
<pre><code class="lang-auto">len_dl = len(dataloader)
running_loss = 0
r_y, r_dy = 0, 0

for batch in dataloader:
    x, y, dy = batch

    optimizer.zero_grad()
                
    # Forward pass
    y_pred, x = model(x)
                
    # Compute gradient
    dy_pred = model.backprop(y_pred) 

    # Compute Loss
    l_y = criterion(y_pred.squeeze(), y)
    l_dy = criterion(dy, dy_pred)

    loss = l_y + l_dy
                
    running_loss += loss.item()
    r_y += l_y.item()
    r_dy += l_dy.item()
                
    # Backward pass
    loss.backward()
    optimizer.step()

running_loss = running_loss/len_dl
r_y = r_y/len_dl
r_dy = r_dy/len_dl
</code></pre>

username: ptrblck
cooked: <p>This error should be raised, since you are attaching tensors to <code>self.zs</code> in each iteration and use them in <code>backprop</code>.<br>
This would let the computation graph grow until you would be running out of memory.<br>
I’m not familiar with your use case, but you could reinitialize the list at the beginning of the <code>forward</code> method via <code>self.zs = []</code>.</p>

username: pascal_notsawo
cooked: <p>You were right, I changed forward like this:</p>
<pre><code class="lang-auto">def forward(self, x):
    zs = []
    for linear_layer in self.net :
        x, z = linear_layer(x)
        zs.append(z)
    y = x
    return y, zs 
</code></pre>
<p>And to correct the problem I just made <code>zbar.detach_()</code> as shown <a href="https://stackoverflow.com/questions/48274929/pytorch-runtimeerror-trying-to-backward-through-the-graph-a-second-time-but" rel="nofollow noopener">here</a>.</p>
<pre><code class="lang-auto">def backprop(self, y, zs):
    zbar = torch.ones_like(y)
    for l in range(len(self.net) - 1, 0, -1):
        linear_layer = self.net[l] 
        zbar =  torch.matmul(zbar, linear_layer.linear.weight) * linear_layer.deriv_activation_function(zs[l-1])
        zbar.detach_() # eq. zbar = zbar.detach()
    linear_layer = self.net[0]
    zbar = torch.matmul(zbar, linear_layer.linear.weight)
    zbar.detach_() # eq. zbar = zbar.detach()
    xbar = zbar
    return xbar 
</code></pre>

username: armarion
cooked: <p>I’d be interested in seeing the pyTorch implementation of differential ML. Do you have a working example you could share?</p>
<p>Thanks!</p>

username: pascal_notsawo
cooked: <p>I did something <a href="https://github.com/Tikquuss/lwd" rel="noopener nofollow ugc">here</a> two years ago. It can be a good start.</p>