I am going to estimate the benchmark on every op in PyTorch in C++. I am going to write a function like below:
# Python
graph = somefunction(model) # <---- Get the computation graph
total_graph_cost = predictCost(graph)

//C++
//wrapper.cc
//use pybind to expose the C++ function to Python.
//...

//Predict function
Cost PredictFunction(Graph graph) {}



Is there any method in PyTorch to get the computation graph and pass it to C++? (Or totally C++ is acceptable).


How could I access the node in the graph? For example, How could I extract the conv2d node’s input tensor shape?


I know that using TorchScript or torch.fx will get the graph. But there are few questions:



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?
How could I access the data-structure in C++? Since this IR may be written in Python (torch.fx case).