Softplus derivative

username: googlebot
cooked: <p>Looking at pytorch sources, I see two implementations:</p>
<p><a href="https://github.com/pytorch/pytorch/blob/994f72ee3e4ae94d73f695613174c9a0a23e8e76/caffe2/operators/softplus_op.cu" rel="nofollow noopener">caffe</a> (no longer used?)</p>
<pre><code class="lang-auto">const float nexpY = exp(-Y[i]);
dX[i] = dY[i] * (1 - nexpY);
</code></pre>
<p><a href="https://github.com/pytorch/pytorch/blob/c177d25edf24a572c09d2f2173f776e015d18d95/aten/src/ATen/native/cuda/Activation.cu" rel="nofollow noopener">aten</a> (a=X, b=Y)</p>
<pre><code class="lang-auto">scalar_t z = std::exp(b * beta);
return (b * beta) &gt; threshold ? a : a * (z - scalar_t(1.)) / z;
</code></pre>
<p>below threshold case is equiv. to caffe:</p>
<pre><code class="lang-auto"> dX = dY * (exp(Y)-1)/exp(Y) = dY * (1 - exp(-Y))
</code></pre>
<p>Notice that caffe version doesn’t need the X tensor. So, are there any problems with using that? (in the most common case, with beta=1, and hessian not needed). Is caffe implementation built-in/exported to python?</p>

username: tom
cooked: <p>To be honest, I don’t think <code>self</code> (the input to the forward) actually used:<br>
</p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/c177d25edf24a572c09d2f2173f776e015d18d95/aten/src/ATen/native/Activation.cpp#L354" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/c177d25edf24a572c09d2f2173f776e015d18d95/aten/src/ATen/native/Activation.cpp#L354" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/c177d25edf24a572c09d2f2173f776e015d18d95/aten/src/ATen/native/Activation.cpp#L354</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="344" style="counter-reset: li-counter 343 ;">
<li>    const Tensor&amp; grad_output,</li>
<li>    const Tensor&amp; self,</li>
<li>    Scalar beta,</li>
<li>    Scalar threshold,</li>
<li>    const Tensor&amp; output) {</li>
<li>  auto iter = TensorIterator::binary_op(grad_input, grad_output, output);</li>
<li>  softplus_backward_stub(iter.device_type(), iter, beta, threshold);</li>
<li>  return grad_input;</li>
<li>}</li>
<li>
</li><li class="selected">Tensor softplus_backward(</li>
<li>    const Tensor&amp; grad_output,</li>
<li>    const Tensor&amp; self,</li>
<li>    Scalar beta,</li>
<li>    Scalar threshold,</li>
<li>    const Tensor&amp; output) {</li>
<li>  Tensor grad_input;</li>
<li>  auto iter = TensorIterator::binary_op(grad_input, grad_output, output);</li>
<li>  softplus_backward_stub(iter.device_type(), iter, beta, threshold);</li>
<li>  return iter.output();</li>
<li>}</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


username: googlebot
cooked: <p>My bad, “a” above binds to grad_output. So, I presume capturing of “self” qualifies as a bug then.</p>

username: googlebot
cooked: <p>On second look, it seems it is done for softplus_double_backward. Interesting tradeoff, simpler (I assume) formula there vs extra memory for a regular case…</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/3f0120edb476af58c3ddf862fc4fde14b73bf927/tools/autograd/templates/Functions.cpp#L1176" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/3f0120edb476af58c3ddf862fc4fde14b73bf927/tools/autograd/templates/Functions.cpp#L1176" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/3f0120edb476af58c3ddf862fc4fde14b73bf927/tools/autograd/templates/Functions.cpp#L1176</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="1166" style="counter-reset: li-counter 1165 ;">
<li>}</li>
<li>
</li><li>Tensor soft_margin_loss_double_backward_grad_output(const Tensor &amp; grad, const Tensor &amp; grad_output, const Tensor &amp; input, const Tensor &amp; target, int64_t reduction) {</li>
<li>  if (reduction == at::Reduction::None) {</li>
<li>    return soft_margin_loss_backward(grad, input, target, reduction);</li>
<li>  }</li>
<li>  auto r = soft_margin_loss_backward(ones_like(grad_output), input, target, reduction);</li>
<li>  return (r * grad).sum();</li>
<li>}</li>
<li>
</li><li class="selected">Tensor softplus_double_backward(const Tensor &amp; grad, const Tensor &amp; input, Scalar beta, Scalar threshold) {</li>
<li>  auto x = (input * beta);</li>
<li>  return sigmoid_backward(grad, x.sigmoid()) * (x &lt; threshold).type_as(grad) * beta;</li>
<li>}</li>
<li>
</li><li>
</li><li>// NOTE [ as_strided Backward and layout-aware/agnostic autograd ]</li>
<li>//</li>
<li>// `storage_offset` is ignored for simplicity in this note. If you just want the</li>
<li>// full algorithm without explanation, scroll down to bottom of this note.</li>
<li>//</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


username: tom
cooked: <p>So one thing one could see if the JIT decomposes this more nicely (I don’t know, but it could be, as that has it’s own differentiation scheme and eliminates unneeded variables).</p>

username: googlebot
cooked: <p>Nope. Here is my test script and alternative version.</p>
<pre><code class="lang-auto">from torch import *
from torch import jit
from torch.nn import functional as F
import gc

#no effect on results
#torch._C._jit_set_profiling_mode(False)
#torch._C._jit_set_profiling_executor(False)

@jit.script
def jit_softplus(x):
	return F.softplus(x)

@jit.script
def fused_softplus(x):
	return x.exp().add(1.0).log()

class LowMemSoftplus(torch.autograd.Function):
	@staticmethod
	def forward(ctx, x):
		y = F.softplus(x)
		ctx.save_for_backward(y)
		return y
	
	@staticmethod
	def backward(ctx, dLoss_dOutput):
		softplus_of_x, = ctx.saved_tensors
		dOutput_dInput = softplus_of_x.neg().exp_()
		torch.sub(1.0, dOutput_dInput, out=dOutput_dInput)
		return dLoss_dOutput * dOutput_dInput

for title, f in {
	"nograd": F.softplus, "grad": F.softplus,
	"jitgrad": jit_softplus, "jitfuse": fused_softplus,
	"custom": LowMemSoftplus.apply,
}.items():
	mem0 = torch.cuda.memory_allocated()
	
	p = torch.ones(5000, device="cuda").requires_grad_(title!="nograd") #this leaf tensor is kept alive due to .grad
	x = p[:, None] @ p[None, :] #non leaf tensor with 25m elements
	
	y = f(x)
	torch.cuda.synchronize()
	
	del x,f
	gc.collect()
	print(title,":",torch.cuda.memory_allocated() - mem0)
	
	del y,p
	gc.collect()
	assert torch.cuda.memory_allocated() == mem0
</code></pre>
<p>output</p>
<pre><code class="lang-auto">nograd : 100683776
grad : 201347072
jitgrad : 201347072
jitfuse : 302010368
custom : 100683776
</code></pre>