Example of creating dot file of libtorch computation graph in c++

torch_autograd_viz.cpp

//
// Created by duane on 10/2/20.
//

#include <torch/torch.h>

int main(int arg, char *argv[]){

    auto x = torch::randn(3, torch::requires_grad());
    auto z = torch::randn(3, torch::requires_grad());
    auto y = x * 2;
    y = y * 2;
    y = y * z;
    y += z;
    y = y * 2;

    std::cout << "digraph G {" << std::endl;

    std::deque<std::shared_ptr<torch::autograd::Node>> nodes;
    nodes.push_front(y.grad_fn());
    while (!nodes.empty()){
        std::cout << "\"" << nodes.back() << "\"" << " [label=\"" << nodes.back()->name() << "\"]" << std::endl;
        for (const auto &edge : nodes.back()->next_edges()) {
            nodes.push_front(edge.function);
            std::cout << "\"" << nodes.back() << "\"" << " -> " << "\"" << edge.function << "\"" << std::endl;
        }
        nodes.pop_back();
    }

    std::cout << "}" << std::endl;

}