How to split backward process wrt each layer of neural network?

zazzyy: <p>Hi everyone,</p>
<p>I’m working on a project that requires me to have access to each step of backward propagation during the training process. Say I have a 10 layer fully connected neural net (input-&gt;fc1-&gt;fc2-&gt;…-&gt;fc10-&gt;output), and during the backward process I want something like output.backward()-&gt;fc10.backward()-&gt;fc9-&gt;backward()-&gt;…-&gt;fc1.backward() in separate steps so that I can get gradient of each layer and check how much time it costs for computing gradient of each layer.</p>
<p>However, for now, when I call <code>loss.backward()</code> using this <code>backward</code> function (<a href="https://github.com/pytorch/pytorch/blob/master/torch/autograd/" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/master/torch/autograd/</a><strong>init</strong>.py#L46 ), then only the <code>loss</code> variable is contained in the <code>variables</code>, and pushed into the <code>execution engine</code>.</p>
<p>How can I get access to gradient computation process of each parameters (layers) in the network? I really want something like</p>
<pre><code class="lang-auto">for layers in reversed(network):
    layer_grad=layers.backward()
    # and at here I can check the time cost of gradient computing of a single layer
   layers.update(layer_grad) # based on certain optimizer
</code></pre>
<p>Any information will be appreciate.</p>

smth: <p>you can use backward hooks for:</p>
<ul>
<li>getting access to the gradient of each of params</li>
<li>checking time taken for computing each layer</li>
</ul>
<p><a href="http://pytorch.org/docs/master/nn.html?highlight=hook#torch.nn.Module.register_backward_hook" class="onebox" target="_blank">http://pytorch.org/docs/master/nn.html?highlight=hook#torch.nn.Module.register_backward_hook</a><br>
<a href="http://pytorch.org/docs/master/autograd.html?highlight=hook#torch.autograd.Variable.register_hook" class="onebox" target="_blank">http://pytorch.org/docs/master/autograd.html?highlight=hook#torch.autograd.Variable.register_hook</a></p>

zazzyy: <p>Hi <a class="mention" href="/u/smth">@smth</a>,</p>
<p>Thanks a lot for the pointers you provide. I also have other questions related, hope you can provide some information.</p>
<p>My project involve solving some straggler (slow worker) issues in distributed cluster. So, I want to “skip” <code>backward</code> calculation (gradient computing) for some layers in the network.</p>
<p>Say if I run the training code on a certain node in a cluster, and after I do the backward process at a certain layer(layer_10.backward -&gt; layer_9.backward -&gt; layer_8.backward), at this time I decided this node is too slow and I want to just simply skip calling <code>backward</code> step of remaining layers to avoid more time costs (e.g. by simply assign a certain value to gradients of those remaining layers rather than actually calculate the gradients) .</p>
<p>Is this possible in current pytorch API? Or I need to customize modules (say convolutional layer)?</p>

smth: <p>What you are asking for is not strictly possible without writing custom autograd.Function functions to insert these custom nodes. Either that, or you do some book-keeping and manage the graph yourself.</p>
<p>For example:</p>
<pre><code class="lang-auto">model = nn.ModuleList(nn.Linear(100, 200), nn.ReLU(), nn.Linear(200, 300))
x = Variable(torch.randn(10, 100), requires_grad=True)

def model_forward(model, x):
    for m in model:
        x = model(x)
        x.detach_()
    return x

def model_backward(model, grad_output):
    for m in reversed(model):
        if TOO_LATE:
            return # shortcut outside of backward
        grad_output = m.backward(grad_output)
    return grad_output        
</code></pre>
<p>Something of this order. It’s super hacky, and you have to do all model book-keeping yourself.</p>

zazzyy: <p>Hi <a class="mention" href="/u/smth">@smth</a>,</p>
<p>I tried this method you provided with following code when defining my customized module list (a simple LeNet example here) and <code>forward</code>, <code>backward</code> operation:</p>
<pre><code class="lang-auto"> class LeNetLayerSplit(nn.Module):
    def __init__(self):
        super(LeNetLayerSplit, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)
        self.ceriation = nn.CrossEntropyLoss()
        self.module_list_0 = nn.ModuleList([self.conv1, nn.MaxPool2d(2, stride=2), nn.ReLU(), 
        					self.conv2, nn.MaxPool2d(2, stride=2), nn.ReLU()])
       	self.module_list_1 = nn.ModuleList([self.fc1, self.fc2])

       	self._name = "LeNet_layer_split"

    def forward(self, x, target):
    	for sub_module in self.module_list_0:
    		x = sub_module(x)
    		x.detach_()
    	x = x.view(-1, 4*4*50)
    	for sub_module in self.module_list_1:
    		x = sub_module(x)
    		x.detach_()
    	loss = self.ceriation(x, target)
    	return x, loss

    def backward(self, grad_output):
    	for m in reversed(self.module_list_1):
    		grad_output = m.backward(grad_output)
    	grad_output.view(-1, 50, 4, 4)
    	for n in reversed(self.module_list_0):
    		grad_output = n.backward(grad_output)
    	return grad_output
</code></pre>
<p>When calling this model, I used the following code:</p>
<pre><code class="lang-auto">def build_model(self):
        self.network = LeNetLayerSplit()
        # this is only used for test
        self.optimizer = torch.optim.SGD(self.network.parameters(), lr=self.lr, momentum=self.momentum)

def train(self, train_loader=None):
        self.network.train()
        # iterate of epochs
        for i in range(self.max_num_epochs):            
            for batch_idx, (data, y_batch) in enumerate(train_loader):
                iter_start_time = time.time()
                data, target = Variable(data, requires_grad=True), Variable(y_batch)
                self.optimizer.zero_grad()
                logits, loss = self.network(data, target)
                print("Trial Loss: {}".format(loss.data[0]))
 
                print("Start Backward Prop Process: ")
                loss.backward()
</code></pre>
<p>But I get the error of <code>RuntimeError: there are no graph nodes that require computing gradients</code>. I guess I call the <code>backward</code> function in a wrong way, and simple search returns no related issue.<br>
But when I read the original code of <code>autograd</code> variable, I found that in this line <a href="https://github.com/pytorch/pytorch/blob/master/torch/autograd/variable.py#L235" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/master/torch/autograd/variable.py#L235</a> results generated during the forward process are set to <code>requires_grad=False</code> when calling <code>detach_</code>, is this issue caused by that? If so, how can I solve it?  Please provide me more details about this.</p>
<p>Thanks a lot!</p>

smth: <p><a class="mention" href="/u/zazzyy">@zazzyy</a> if you call <code>detach_()</code> then there is no node in the graph that requires_grad=True, so autograd is complaining that it has no work to do.</p>
<p>What you might want to do is (maybe), instead of <code>x.detach_()</code>, call <code>x = Variable(x.data, requires_grad=True)</code> (or some form of this, that will compute gradients).</p>

zazzyy: <p>Hi, <a class="mention" href="/u/smth">@smth</a>  thanks a lot for your response. Based on your suggestions, I tried the following things bellow:</p>
<ol>
<li>I tried to remove every <code>x.detach_()</code> in the foregoing code and simply called <code>loss.backward()</code>, then the model works as the normal condition and my customized <code>backward()</code> function are not called.</li>
</ol>
<pre><code class="lang-auto">def backward(self, grad_output):
    	for m in reversed(self.module_list_1):
    		grad_output = m.backward(grad_output)
    	grad_output.view(-1, 50, 4, 4)
    	for n in reversed(self.module_list_0):
    		grad_output = n.backward(grad_output)
    	return grad_output
</code></pre>
<ol start="2">
<li>When I tried to call the customized <code>backward()</code> function in this manner, <code>self.network.backward(grad_output=${RANDOM_VARIABLE})</code>, then I got this error:</li>
</ol>
<pre><code class="lang-auto">    grad_output = m.backward(grad_output)
  File "/home/usr/anaconda2/lib/python2.7/site-packages/torch/nn/modules/module.py", line 262, in __getattr__
    type(self).__name__, name))
AttributeError: 'Linear' object has no attribute 'backward'
</code></pre>
<p>Dose this mean I need to overwrite the corresponding <code>nn.Module</code> by myself with a <code>backward</code> function, or there is a way that the current <code>Module</code> will do that for me? So, in general, I only want to call the backward process layer by layer, after getting the gradient of each layer, I just want to do a simple check (e.g. the value of a certain variable) and to determine if I need to do or just skip the remaining backward process.</p>

smth: <p>in your backward function, you are calling <code>.backward</code> on <code>m</code> which is the Module. what you need to do is to call <code>backward</code> on the Variable that is the output from the module.</p>

zazzyy: <p>Hi, <a class="mention" href="/u/smth">@smth</a>. Thanks a lot for this pointer.</p>
<p>I already tried what you suggest that call the <code>.backward()</code> on the output Variable that is output from the module.  The case is, if I do the normal <code>forward()</code> manner like this way:</p>
<pre><code class="lang-auto">    def forward(self, x, target):
    	for sub_module in self.module_list_0:
    		x = sub_module(x)
    	x = x.view(-1, 4*4*50)
    	for sub_module in self.module_list_1:
    		x = sub_module(x)
    	loss = self.ceriation(x, target)
    	return x, loss
</code></pre>
<p>Then once I call the <code>variable.backward()</code> on the last variable, say <code>loss.backward()</code> then the whole backward process will be executed. But if I call anything like <code>x = Variable(x.data, requires_grad=True)</code> after each forward step as you mentioned, then it seems no grad will be calculated if I check <code>param.grad</code> in <code>module.parameters()</code>.<br>
After checking this topic <a href="https://discuss.pytorch.org/t/assign-manual-assigned-grad-output/2750">Assign manual assigned "grad_output"</a>, it seems call <code>variable.backward(grad_output)</code> can be helpful, but when I calling <code>loss.backward(grad_output)</code> under normal forward manner, the behavior is nothing different from <code>loss.backward()</code>.<br>
What tricks do I need in the forward/backward process to achieve executing backward process layer by layer (e.g. get output from last layer and use it to do backward for next layer, like doing the chain rule manually )?</p>
<p>Thanks a lot!</p>

smth: <p>here’s a more precise and fuller example. What you are doing in my example is to completely avoid autograd’s automatic backward computation and manually reverse-computing the backward graph.</p>
<p><strong>For anyone  coming here with a search, my solution is a hack, it is not good practice. it is given as an illustration just to showcase to <a class="mention" href="/u/zazzyy">@zazzyy</a> how to shortcut these things</strong></p>
<pre><code class="lang-python">import torch
import torch.nn as nn
from torch.autograd import Variable

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layers = nn.ModuleList([
            nn.Linear(10, 10),
            nn.Linear(10, 10),
            nn.Linear(10, 10),
            nn.Linear(10, 10),
        ])

    def forward(self, x):
        self.output = []
        self.input = []
        for layer in self.layers:
            # detach from previous history
            x = Variable(x.data, requires_grad=True)
            self.input.append(x)

            # compute output
            x = layer(x)

            # add to list of outputs
            self.output.append(x)
        return x

    def backward(self, g):
        for i, output in reversed(list(enumerate(self.output))):
            if i == (len(self.output) - 1):
                # for last node, use g
                output.backward(g)
            else:
                output.backward(self.input[i+1].grad.data)
                print(i, self.input[i+1].grad.data.sum())

model = Net()
inp = Variable(torch.randn(4, 10))
output = model(inp)
gradients = torch.randn(*output.size())
model.backward(gradients)
</code></pre>

zazzyy: <p>Hi, <a class="mention" href="/u/smth">@smth</a>. I deeply appreciate this working snippet, it is basically what I asked.</p>
<p>I have one further question about this. It seems the <code>grad_output</code> we provide in this line</p>
<pre><code class="lang-auto">gradients = torch.randn(*output.size())
</code></pre>
<p>are never used to compute anything during the backward process since when I tried to print every <code>self.output[i+1].grad.data</code>, they’re all equal to the random gradient generated by <code>torch.randn(*output.size())</code>. And everything in <code>self.output</code> has <code>grad_fn=None</code>. That make sense because as you mentioned this hack “ompletely avoid autograd’s automatic backward computation and manually reverse-computing the backward graph”.</p>
<p>My question is if I want to make this net works “normally”, do I need to manually handle the entire backward process (e.g. write all backward functions for all layers, manually compute gradients, and etc.)? Is there any way to make this less hacky or any function (e.g. backward function of Conv layer) to borrow?</p>

smth: <p>I’m sorry the example had a bug. I’ve fixed it now (see my example again) and it correctly computes gradients.</p>

samarth-robo: <p><a class="mention" href="/u/smth">@smth</a> can you use hooks to measure the time taken for the backward pass of a module?</p>

Smart_Zhang: <p>Is there any examples of measuring the backward propagation time of a layer by using hooks?</p>

maralm: <p>May I know why it is not good practice? Are there other ways of splitting the layers?</p>