How to get Operator from torch::jit::module

username: Stonepia
cooked: <p>Hi,</p>
<p>How could I get the Operator from torch::jit::module?<br>
I want to traverse the whole module and estimate the costs on every op (since I want to get custom costs result, so using <code>torch.autograd.profiler</code> is not that applicable).</p>
<pre><code class="lang-auto">// some_cpp_code.cpp
  &lt;...&gt;
  m.def("get_cost", [](const torch::jit::Module&amp; mod) {
    auto forward_method = mod.get_method("forward");
    assert(forward_method.graph());
    Costs cost;

    // Traverse the graph, is there any more intuitive way? 
    for (Node * n : forward_method.graph()-&gt;nodes()) {
      auto kind = n-&gt;kind();
      // Calculate the cost
      switch (kind){
          case aten::add:
                cost += my_add_cost( n-&gt;inputs(), n-&gt;outputs());
          case aten::conv2d:
                cost += my_conv2d_cost(n-&gt;inputs(), n-&gt;outputs());
       //...
    }
  });

</code></pre>
<p>The above way could (possibly) work, but it makes the code messy and actually I don’t need that much information. I just need the <code>Operator</code> and <code>inputs</code> and <code>outputs</code>. The ideal function call would be something like :</p>
<pre><code class="lang-auto">Costs my_cost(operator, inputs, outputs){
// Using a map to register correct cost function
}

// cost function
Costs my_conv2d_cost(operator, inputs, outputs){
// Get other information about this op
auto padding = inputs[4]; // Is this right?
}
</code></pre>
<p>I have 3 questions:</p>
<ol>
<li>Is PyTorch offers any way to traverse the whole module and get every op? The above code would have to use a <code>switch</code> case to judge every case, this is quite annoying. Also, it will encounter many nodes like <code>prim::constant</code> , and those nodes are kind of not interesting.</li>
<li>Is <a href="https://github.com/pytorch/pytorch/blob/4a0d17ba2d5966cea27dbce527baef62bb50c6f7/torch/csrc/jit/runtime/operator.h#L120-L128" rel="noopener nofollow ugc"><code>getOperation(const Node* node = nullptr)</code></a> function the one I need? It returns the operation, but actually I don’t quite understand how to use it , it doesn’t filter out the <code>prim::constant</code> node as stated before, right?</li>
<li>How could I get the correct inputs, like the code shown of <code>my_conv2d_cost</code>, is this the right way?</li>
</ol>

username: SplitInfinity
cooked: <ol>
<li>No, but you can use <code>Node::maybeOperator</code> to skip nodes that don’t have operators (like <code>prim</code> nodes).</li>
<li>No; <code>getOperation</code> returns an <code>Operation</code>, which is an alias for <code>std::function&lt;void(Stack*)&gt;</code>. This is a callable function that executes the operation, but will not be useful for cost analysis. Consider using <code>getOperator</code> or <code>maybeOperator</code>; these functions return an <code>Operator</code> that has an associated schema (<code>Operator::schema()</code>) that can be used to determine which operator it is (essentially the same as <code>n-&gt;kind()</code> from your example).</li>
<li>Take a look at the functions in <code>jit/ir.h</code>; you will find functions there that help you access node inputs and their types. However, in the case of tensor types, you might not always find complete shape and dtype information (it depends on how the graph is produced and when your analysis code runs).</li>
</ol>

username: Stonepia
cooked: <p><a class="mention" href="/u/splitinfinity">@SplitInfinity</a></p>
<p>Thanks for the reply! Could you give me some hints about this line?</p>
<blockquote>
<p>However, in the case of tensor types, you might not always find complete shape and dtype information (it depends on how the graph is produced and when your analysis code runs).</p>
</blockquote>
<p>I have also found that it is hard to get the info from incomplete tensor ( conv2d’s <code>weight</code> tensor for example). Could you give me some hints on how pytorch handles this case? Rerun the full module seems not a good idea for large module case in my opinion.</p>

username: SplitInfinity
cooked: <p>Unfortunately, there is no way to propagate shape information other than running the model. We are working on ways to let users add this information to graphs directly but this won’t be released anytime soon.</p>

username: Stonepia
cooked: <p>Thanks for your work!</p>
<blockquote>
<p>We are working on ways to let users add this information to graphs directly</p>
</blockquote>
<p>Do you mean the <a href="https://github.com/pytorch/rfcs/blob/rfc-0005/RFC-0005-structured-kernel-definitions.md" rel="noopener nofollow ugc">structured kernel definitions</a>? If so then maybe I could try to refer to existing PRs?</p>

username: SplitInfinity
cooked: <p>No, that is different. There is ongoing work on a tensor DSL but unfortunately there’s no public RFC or something I can share.</p>