Neural network visualization - calculation graph of the model drawn based on torchviz

  Visualizing the network structure is a very useful tool when understanding deep learning models. Today we will introduce how to use torchviz and graphviz to generate network calculation graphs. This approach is particularly useful for readers who wish to delve deeper into the internal details of a network. It should be noted that the generated network structure diagram is generated based on the backpropagation process, so it shows a network structure in reverse order.

The first step is to install graphviz and torchviz libraries

  First, we need to install two Python libraries: torchviz and graphviz. These two libraries work together to help us generate intuitive network structure diagrams. Installing these libraries is as simple as running the following code.

pip install graphviz torchviz

Step 2: Write code to generate calculation graph

  Next, we need to write a piece of code to generate the calculation graph. This code will use the previously installed libraries to create and display the network structure.

import torch
import torch.nn as nn
from torchviz import make_dot

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(5, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleModel()
input_tensor = torch.randn(1, 10)  # 假设模型输入是一个 10 维的向量


output = model(input_tensor)
graph = make_dot(output, params=dict(model.named_parameters()))
graph.render('model_graph', format='png')  # 保存为 PNG 图像

  After running the above code, you may encounter an error:

  The error message is that some executable commands in graphviz have not been added to the system path. This means that although we have installed the relevant Python libraries, the graphviz software itself has not been installed. The solution to this problem is to download and install the graphviz software and make sure its executable is added to the system environment variables. The graphviz software can be downloaded from the following link:https://www2.graphviz.org/Packages/stable/windows/10/cmake/Release/x64/

Step 3. Install graphviz software

  Download and run the installer. During the installation process, make sure to select the option to add graphviz to the system environment variables. This step is very important as it allows your code to properly call graphviz functionality.

  When installing, remember to add system environment variables:

  After the installation is complete, if the code still cannot run, try restarting the computer. Doing this often resolves issues related to updating environment variables. After installing and configuring all required components, you should be able to successfully generate and view computational graphs. An example of the final generated calculation graph is as follows:

  Computational graph is a directed graph that mainly describes the relationship between operations and variables in the network.

  At the top of the figure, you can see the weights fc1.weight and biases fc1.bias of the network layer, which represent the first fully connected layer in the network ( Usually represented by fc) parameters. The weight matrix has dimensions (5, 10) and the bias vector has dimension (5). This means that this layer will take a 10-dimensional input and output a 5-dimensional vector.

  The nodes in the middle represent different operations in the backpropagation process, such as TBackward or AddmmBackward. These are automatic gradient calculation operations in PyTorch. Used to update network weights during training.

  fc2.weight and fc2.bias represent the parameters of the second fully connected layer, with the weight matrix dimension being (2, 5) and the bias vector dimension being (2). This shows that the second layer converts 5-dimensional input into 2-dimensional output.

  The bottom one (1, 2) represents the final output of the network, which is a 2-dimensional vector.