How to get the computation graph and pass it to C++?

Stonepia: <p>I am going to estimate the benchmark on every op in PyTorch in C++. I am going to write a function like below:</p>
<pre><code class="lang-auto"># Python
graph = somefunction(model) # &lt;---- Get the computation graph
total_graph_cost = predictCost(graph)
</code></pre>
<pre><code class="lang-auto">//C++
//wrapper.cc
//use pybind to expose the C++ function to Python.
//...

//Predict function
Cost PredictFunction(Graph graph) {}
</code></pre>
<ol>
<li>
<p>Is there any method in PyTorch to get the computation graph and pass it to C++? (Or totally C++ is acceptable).</p>
</li>
<li>
<p>How could I access the node in the graph? For example, How could I extract the <code>conv2d</code> node’s input tensor shape?</p>
</li>
<li>
<p>I know that using TorchScript or torch.fx will get the graph. But there are few questions:</p>
</li>
</ol>
<ul>
<li>These two methods get two different IR. How could I access the IR in C++? For example, torch.fx totally access it by python, is C++ code available to access this IR?</li>
<li>How could I access the data-structure in C++? Since this IR may be written in Python (torch.fx case).</li>
</ul>

James_Reed: <p>Hi <a class="mention" href="/u/stonepia">@Stonepia</a>,</p>
<p>Thanks for the question!</p>
<p>First, I’d like to understand better what you intend to do here. Are you trying to calculate a quantity like FLOPs for the network? One way to do this is with the profiler: <a href="https://pytorch.org/docs/stable/autograd.html#profiler" class="inline-onebox" rel="noopener nofollow ugc">Automatic differentiation package - torch.autograd — PyTorch 1.8.0 documentation</a>. It seems to do much of what you’re trying to achieve, however its interface is not in C++ but rather Python.</p>
<p>To answer your direct questions:</p>
<ol>
<li>I think the most out-of-the-box way to do this is to use TorchScript. For example, you can use tracing and a Pybind API that takes <code>torch::jit::Module</code> like so:</li>
</ol>
<pre><code class="lang-auto"># some_cpp_code.cpp
  &lt;...&gt;
  m.def("_test_fn", [](const torch::jit::Module&amp; mod) {
    auto forward_method = mod.get_method("forward");
    assert(forward_method.graph());
    for (Node * n : forward_method.graph()-&gt;nodes()) {
      n-&gt;dump();
    }
  });
</code></pre>
<pre><code class="lang-auto"># Python code for testing
import torch
import torchvision.models as models

# Construct and trace ResNet18 for testing
rn18 = models.resnet18()
traced_rn18 = torch.jit.trace(rn18, (torch.randn(5, 3, 224, 224),))

# Inline method calls. This allows us to more easily traverse the graph
# of operations
torch._C._jit_pass_inline(traced_rn18.graph)
# Pass our module into the custom C++ function. Note that we use `._c` to
# retrieve the C++ version, not the Python version
_test_fn(traced_rn18._c)
</code></pre>
<p>Running that, you will see print-outs of all the nodes in the model.</p>
<ol start="2">
<li>When using tracing, the shapes of values are recorded, you can look at them by using <code>type()</code> on each node. However, note that this is specific to the invocation that was used to record the graph</li>
<li>
<code>torch.fx</code> is exclusively Python. You can define a C++ structure isomorphic to the FX graph, but it will probably be more work. We don’t have plans to generalize FX to C++ at this point</li>
</ol>
<p>James</p>

Stonepia: <p><a class="mention" href="/u/james_reed">@James_Reed</a>  Thank you very much for your answer, you saved my life!</p>
<p>Yes, I want to do something similar like profiler. However, the profiler requires to truly run the model, while I just want to get a rough prediction of the cost based on its input/output size (just as tf’s <a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/grappler/costs/analytical_cost_estimator.h" rel="noopener nofollow ugc">grappler</a> does). Then using torch’s profiler along seems not working.</p>
<p>I want to implement a custom cost estimator based on every op , and finally get the final model cost.</p>
<p>In this way, I want to get the graph and use the code like your suggestion:</p>
<pre><code class="lang-auto">Costs totalcost ;
for (Node * n : forward_method.graph()-&gt;nodes()) {
     totalcost += predictCost(n);
}
</code></pre>
<p>Thank you again, your reply really helps!</p>

Sujoy_Saraswati: <p>Is there any plan to introduce an FX graph in C++ eventually?</p>