Separating out forward & backward passes as tensor operations

username: nirandaperera
cooked: <p>Hi,</p>
<p>I am trying to separate out forward &amp; backward passes of <code>nn.module</code> as tensor operations (I am working on a peculiar requirement here <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>
<p>The interfaces look like this.</p>
<pre data-code-wrap="python"><code class="lang-python">class LinearForward:
    def __init__(self, in_features: int, out_features: int, bias=True, device=torch.device("cpu")):
    ...
    def __call__(self, input: Tensor) -&gt; Tensor:
        self.input = input
        output = input.matmul(self.weight.t())
        if self.bias is not None:
            output += self.bias
        return output

class LinearBackward:
    def __init__(self, forward: LinearForward, device=torch.device("cpu")):
    ...
    def __call__(self, grad_output: Tensor) -&gt; Tuple[Tensor, Tensor, Any]:
        grad_input = grad_output.matmul(self.weight)
        grad_weight = grad_output.t().matmul(self.forward.get_input())
        grad_bias = None
        if self.bias is not None:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias
</code></pre>
<p>And I have tested the implementation against <code>nn.Linear</code> implementations and observed the following in cuda devices.</p>
<ol>
<li>my tensor implementation is about 10-15% slower than the <code>nn.Linear</code></li>
<li>In both cases, almost all the time is spent on the backward pass, while the forward takes &lt;1% of the training time (at least for the Linear layer)</li>
</ol>
<p>What would be the reason for this? Is it because my forward and backward classes are written in python and each <code>__call__</code> method this getting synchronized with the python GIL (Like in the forward and backward hook implementations in <code>nn.Module</code>s)?</p>
<p>Look forward to hearing from the community.</p>
<p>PS:<br>
I have followed these threads previously.</p><aside class="quote" data-post="7" data-topic="2513">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ngimel/48/14766_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/measuring-gpu-tensor-operation-speed/2513/7">Measuring GPU tensor operation speed</a> 
  </div>
  <blockquote>
    The overhead on second iteration is garbage collection/reference counting not keeping up, and caching allocator allocating second y tensor. If you add y=None in the timing loop, you’d be able to reuse y allocation, and only the first iteration would show additional overhead.
  </blockquote>
</aside>
<aside class="quote quote-modified" data-post="1" data-topic="53922">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fruitvinegar/48/39246_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/analyzing-hook-drastically-decreases-performance/53922">Analyzing Hook drastically decreases performance</a> <a class="badge-category__wrapper " href="/c/deployment/14"><span data-category-id="14" data-drop-close="true" class="badge-category " title="A category of posts focused on production usage of PyTorch. Mobile deployment is out of scope for this category (for now…  )"><span class="badge-category__name">deployment</span></span></a>
  </div>
  <blockquote>
    Hi. I’m using a PyTorch Hook to analyzing intermediate feature map’s statistics with a large-scale model training. 
Also I’m using multi-GPU to accelerate the training process.(with torch.nn.DataParallel) 
The problem is, GPU utilization is drastically drops down when my Hook is executing(I compared GPU utilization with when I don’t activated my Hook - naive training without analyzing) 
I know this is Indispensable situation because my analyzing Hook is quite high-cost function, but all my multi…
  </blockquote>
</aside>


username: albanD
cooked: <p>Hi,</p>
<p>Do you use <code>torch.cuda.synchronize()</code> properly when doing your timing? That sounds weird that the forward is &lt;1%.</p>

username: nirandaperera
cooked: <p>Hey <a class="mention" href="/u/alband">@albanD</a>,<br>
Oh, you are correct. That was my mistake.<br>
Could you help clarifying (1)?</p>

username: albanD
cooked: <p>The main reason I can think of is that when you use autograd, we only compute gradients for the things that require gradients. In particular, if the input does not require gradients, then autograd won’t compute <code>grad_input</code> which might explain some perf difference.<br>
Also the function runs purely in cpp which might also explain a few % of difference, especially if you test with small input Tensors.</p>