How to implement a custom activation function in the PyTorch C++ frontend?

tmaric: <p>I’m having difficulties finding documentation that describes extending the C++ frontend online, specifically, implementing a custom activation function. Before digging through the source code I wanted to ask here if anyone here has any information on that. Thanks!</p>

ptrblck: <p><a href="https://pytorch.org/tutorials/advanced/torch_script_custom_classes.html">This tutorial</a> explains how to write custom C++ classes and call them in the Python frontend as well as in libtorch.<br>
Is this what you are looking for?</p>

tmaric: <p>I’m basically looking to understand how to implement an activation function in the C++ frontend, and stay in C++, since I’m using PyTorch in another C++ software. AFAIK TorchScript is used for python integration, right?</p>
<p>I’m digging throught the C++ API source code, and I’ve found this</p>
<pre><code class="lang-auto">class TORCH_API ELUImpl : public torch::nn::Cloneable&lt;ELUImpl&gt; {
 public:
  explicit ELUImpl(const ELUOptions&amp; options_ = {});

  Tensor forward(Tensor input);

  void reset() override;

  /// Pretty prints the `ELU` module into the given `stream`.
  void pretty_print(std::ostream&amp; stream) const override;

  /// The options with which this `Module` was constructed.
  ELUOptions options;
};

/// A `ModuleHolder` subclass for `ELUImpl`.
/// See the documentation for `ELUImpl` class to learn what methods it
/// provides, and examples of how to use `ELU` with `torch::nn::ELUOptions`.
/// See the documentation for `ModuleHolder` to learn about PyTorch's
/// module storage semantics.
TORCH_MODULE(ELU);
</code></pre>
<p>in</p>
<p><code>include/torch/csrc/api/include/torch/nn/modules/activation.h</code></p>
<p>So activation functions seem to be wrapped into function objects that implement <code>forward</code>, but I can’t seem to find the actual implementation. Also, what does <code>TENSOR_MODULE(ELU)</code> do? Register the class somehow somewhere?</p>
<p>I’m also having difficulties finding an actual implementation of the elu funciton.</p>
<p>Specifically,</p>
<pre><code class="lang-auto">namespace detail {
inline Tensor elu(Tensor input, double alpha, bool inplace) {
  if (inplace) {
    return torch::elu_(input, alpha);
  } else {
    return torch::elu(input, alpha);
  }
}
} // namespace detail
</code></pre>
<p>forwards to <code>torch::elu_</code> which is then</p>
<p><code>CAFFE2_API Tensor &amp; elu_out(Tensor &amp; out, const Tensor &amp; self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1);</code></p>
<p>a CAFFE_2API native function?</p>
<p>So, back to my question, I want to implement a Gaussian as an activation function, what do I need to do?</p>
<ol>
<li>Implement the function in <code>detail::</code> ?</li>
<li>Implement the <code>Impl</code> class (function object, module?) ?</li>
<li>“Register” the module?<br>
…</li>
</ol>
<p>I mean, I could try copying what I see, and digging through the code, but it would be nice if there was some information on the design available somewhere</p>

tmaric: <aside class="quote no-group" data-username="tmaric" data-post="3" data-topic="139713">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/f19dbf/48.png" class="avatar"> tmaric:</div>
<blockquote>
<p><code>ModuleHolder</code></p>
</blockquote>
</aside>
<p>I finally found some information in <code>include/torch/csrc/api/include/torch/nn/module.h</code>… reading this now.</p>