Custom activation function is slow

Jow1e: <p>Hi!</p>
<p>I’m new here and i am writing custom activation function with learnable parameters:</p>
<p>How it looks:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/9/c9f86c3f796b6a904ecbc6b47a4fb30536f5c298.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/c9f86c3f796b6a904ecbc6b47a4fb30536f5c298" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/9/c9f86c3f796b6a904ecbc6b47a4fb30536f5c298_2_690x482.jpeg" alt="image" data-base62-sha1="sOIdkDVvToK8H9EPrn7GyjxBKVa" width="690" height="482" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/9/c9f86c3f796b6a904ecbc6b47a4fb30536f5c298_2_690x482.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/c/9/c9f86c3f796b6a904ecbc6b47a4fb30536f5c298_2_1035x723.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/c/9/c9f86c3f796b6a904ecbc6b47a4fb30536f5c298.jpeg 2x" data-dominant-color="EDE7E9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1314×919 85.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The formula:<br>
<em>x[ x &lt; 0 ] = x[ x &lt; 0 ] * alpha_low</em><br>
<em>x[ x &gt; bias ] = x[ x &gt; bias ] * alpha_high + bias</em></p>
<p>I tried several ways implementing it and here is the fastest result:</p>
<pre><code class="lang-auto">class CustomRelu(nn.Module):
	def __init__(self, alpha_low=0.01, alpha_high=0.01, bias=1.):
		super(CustomRelu, self).__init__()
		
		self.prelu_low = nn.PReLU(init=alpha_low)
		self.prelu_high = nn.PReLU(init=alpha_high)
		
		self.bias = bias
	
	def forward(self, x):
		x= -self.prelu_low(x) + self.bias
		x= -self.prelu_high(x) + self.bias
		return x
</code></pre>
<p>But the problem is, its still much slower than built in activation functions (about 5-10 times slower)<br>
Is there a way how to optimize it?</p>

eqy: <p>Is this operation done on CPU or GPU? One potential for optimization on GPU here is writing a custom kernel/extension: <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" class="inline-onebox" rel="noopener nofollow ugc">Custom C++ and CUDA Extensions — PyTorch Tutorials 1.9.0+cu102 documentation</a> to fuse the operations together, because as it is written, the activation function incurs many global memory read/writes. From a quick glance, x is read, negated, written back, read again to add bias, written back, read again to negate, written back, read again to add bias, and finally written back to produce the output. With a single kernel, this could be optimized to a single read/write. However, in addition to requiring a new kernel/CUDA extension to be written you would also need to implement your own backwards kernel for training as the new operation would not be registered in autograd.</p>

Jow1e: <p>Thank you. I’ll try it!</p>

eqy: <p>If this works for you use case, please update this thread with your performance results! It could be a great example for others also trying to optimize similar operators/kernels.</p>