Accessing parameter names of jit traced autograd.Function's

username: iutlu
cooked: <p>Hello, I’ve been trying to convert a model with a custom <code>autograd.Function</code> to coreml.</p>
<pre><code class="lang-auto">class Custom(autograd.Function):
    @staticmethod
    def forward(ctx, input1, input2, param1, param2, param3):
        ...

    @staticmethod
    def backward(ctx, grad_output):
        ...
</code></pre>
<p>When traced, invocations of <code>Custom.apply</code> result in <code>node</code>s of kind <code>prim::PythonOp</code>, with <code>input1</code>, <code>input2</code> as <code>node.inputs()</code>. Different from other types of ops, the scalar parameters <code>param1</code>, <code>param2</code>, <code>param3</code> do not seem to be registered as <code>prim::Constant</code> nodes in the graph (as far as I can see).</p>
<p>I was hoping to monkey-patch <code>InternalTorchIRNode</code> from <a href="https://github.com/apple/coremltools/blob/5.1/coremltools/converters/mil/frontend/torch/internal_graph.py#L114" rel="noopener nofollow ugc">coremltools</a> and handle these parameters separately. Specifically, my thinking was to construct <code>constant</code>-kinded <code>InternalTorchIRNode</code>s <em>as if</em> <code>torch.jit.trace</code> had produced the corresponding <code>prim::Constant</code> nodes, and add them to the <code>InternalTorchIRGraph</code> instance myself, and wiring everything else together.</p>
<p>This is seems pretty doable: I’m readily able to access the values of these parameters via <code>node.scalar_args()</code>. I also wanted to access the <em>names</em> of the parameters (in order to be able to label my coreml nodes), and here’s where things got tricky: I can get a reference to the <code>Custom.apply</code> function via <code>node.pyobj()</code>, however it’s not possible to get a reference to <code>Custom</code> itself from here (getting such a reference would allow me to inspect the signature of <code>Custom.forward</code>, which is what I’m after). <code>node.pyname()</code> is not helpful either, since it’s not possible in general to unambiguously discover the autograd function by the name <code>'Custom'</code> alone.</p>
<p>Looking around for a solution, the following caught my eye:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/v1.10.2/torch/csrc/jit/python/python_ir.cpp#L157">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/v1.10.2/torch/csrc/jit/python/python_ir.cpp#L157" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/v1.10.2/torch/csrc/jit/python/python_ir.cpp#L157" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/v1.10.2/torch/csrc/jit/python/python_ir.cpp#L157</a></h4>



    <pre class="onebox">      <code class="lang-cpp">
        <ol class="start lines" start="147" style="counter-reset: li-counter 146 ;">
            <li>  this-&gt;pyobj = THPObjectPtr(other-&gt;pyobj.get());</li>
            <li>  for (auto&amp; sa : other-&gt;scalar_args) {</li>
            <li>    Py_INCREF(sa.get());</li>
            <li>    this-&gt;scalar_args.emplace_back(sa.get());</li>
            <li>  }</li>
            <li>}</li>
            <li>
            </li>
<li>// recover the autograd.Function instance, if this PythonOp's function</li>
            <li>// was originally SomeFunction.apply</li>
            <li>// used in ONNX for discovering symbolics</li>
            <li class="selected">c10::optional&lt;THPObjectPtr&gt; ConcretePythonOp::autogradFunction() const {</li>
            <li>  pybind11::gil_scoped_acquire gil;</li>
            <li>  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)</li>
            <li>  py::handle obj = const_cast&lt;PyObject*&gt;(pyobj.get());</li>
            <li>
            </li>
<li>  auto r = py::getattr(obj, "__self__", py::none());</li>
            <li>  if (r.is_none())</li>
            <li>    return c10::nullopt;</li>
            <li>
            </li>
<li>  auto apply = py::getattr(r, "apply", py::none());</li>
            <li>  if (apply.is_none())</li>
        </ol>
      </code>
    </pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>It looks like <code>autogradFunction</code> method from <code>ConcretePythonOp</code> is exactly what I’m looking for. It doesn’t seem to be exposed to the Python side, however, which brings me to the question here. Is there a way to access these parameter names, or could <code>autogradFunction</code> be exposed in <code>torch._C.Node</code>?</p>
<p>I don’t really understand everything that’s involved – possibly what I’m asking for is not technically possible – but I thought I’d ask here anyways, thanks</p>

username: iutlu
cooked: <p>Silly me, all I needed was to read on to the implementation <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
(TIL: the <code>__self__</code> attribute)</p>
<pre><code class="lang-auto">from inspect import signature
params = list(signature(node.pyobj().__self__.forward).parameters)
</code></pre>
<p>There it is. Hope it’s useful to someone.</p>