PyTorch is a popular open-source machine learning library for developing deep learning models. It provides a wide range of functions for building complex neural networks. PyTorch defines a computational graph as a Directed Acyclic Graph (DAG) where nodes represent operations (e.g., addition, multiplication, etc.) and edges represent the flow of data between the operations. When defining a PyTorch model, the computational graph is created by defining the forward function of the model. This function takes inputs and applies a sequence of operations to produce the outputs. During the forward pass, PyTorch creates the computational graph by keeping track of the operations and their inputs and outputs.
A model’s calculation is represented by a graph, a type of data structure. A graph is represented in PyTorch by a directed acyclic graph (DAG), where each node denotes a calculation and each edge denotes the data flow between computations. In the backpropagation step, the graph is used to compute gradients and monitor the dependencies between computations.
A computational graph is a graphical representation of a mathematical function or algorithm, where the nodes of the graph represent mathematical operations, and the edges represent the input/output relationships between the operations. In other words, it is a way to visualize the flow of data through a system of mathematical operations.
Computational graphs are widely used in deep learning, where they serve as the backbone of neural networks. In a neural network, each node in the computational graph represents a neuron, and the edges represent the connections between neurons.
Implementations
Install the necessary libraries
sudo apt install graphviz # [Ubuntu] winget install graphviz # [Windows] sudo port install graphviz # [Mac] pip install torchviz
Step 1: Import  the necessary libraries
Python3
import torchimport torch.nn as nnimport torch.optim as optimimport torchvision.datasets as datasetsimport torchvision.transforms as transformsfrom torch.utils.data import DataLoaderfrom PIL import ImageÂ
from torchviz import make_dot |
Step 2: Build the neural network model
Python3
# Define a neural network architectureclass Net(nn.Module):Â Â Â Â def __init__(self):Â Â Â Â Â Â Â Â super(Net, self).__init__()Â Â Â Â Â Â Â Â self.conv1 = nn.Conv2d(3, 16, 3)Â Â Â Â Â Â Â Â self.pool = nn.MaxPool2d(3, 2)Â Â Â Â Â Â Â Â self.fc1 = nn.Linear(238, 120)Â Â Â Â Â Â Â Â self.fc2 = nn.Linear(120,1)Â Â Â Â Â Â Â Â self.fc3 = nn.Sigmoid()Â
    def forward(self, x):        x = self.pool(torch.relu(self.conv1(x)))        x = x.view(-1, 238)        x = torch.relu(self.fc1(x))        x = torch.relu(self.fc2(x))        x = torch.relu(self.fc3(x))        return x     # Move the model to the GPUdevice = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")net = Net().to(device) |
Step 3: Load the image datasetÂ
Python3
# Load the imageimg = Image.open("Ganesh.jpg")Â
# Transformationtransform = transforms.Compose([    transforms.Resize((720, 480)), # Resize the image to 720x480    transforms.ToTensor(),    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])Â
img = transform(img)# Reshape the 4d tensorimg = img.view(1, 3, 720, 480)Â
# Move to GPUimg = img.to(device)img.shape |
Output:
torch.Size([1, 3, 720, 480])
Step 4: Process the first convolutional layer
Python3
# first convolutional layerconv1 = net.conv1print('Convolution Layer :',conv1)Â
# APPLY THE FIRST Convolutional layer to the imagey1 = conv1(img)print('Output Shape :',y1.shape)Â
# Computational Graphmake_dot(y1, params=dict(net.named_parameters())) |
Output:
Convolution Layer : Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1)) Output Shape : torch.Size([1, 16, 718, 478])
Computational Graph
Step 5: Process the Second Pooling layer
Python3
# Pooling layerpool = net.poolprint('pooling Layer :',pool)Â
# APPLY THE Pooling layer to the# output of FIRST Convolutional layery2 = pool(y1)print('Output Shape :',y2.shape)Â
# Computational Graphmake_dot(y2, params=dict(net.named_parameters())) |
Output:
pooling Layer : MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False) Output Shape : torch.Size([1, 16, 358, 238])
Computational Graph
Step 6: Process the first Linear layer
Python3
# first Linear layerfc = net.fc1print('First Linear Layer :',fc)Â
# APPLY THE Linear layer to the output y2y3 = fc(y2)print('Output Shape :',y3.shape)Â
# Computational Graphmake_dot(y3, params=dict(net.named_parameters())) |
Output:
First Linear Layer : Linear(in_features=238, out_features=120, bias=True) Output Shape : torch.Size([1, 16, 358, 120])
Computational Graph
Step 7: Process the second Linear layer
Python3
# Second Linear layerfc_2 = net.fc2print('Second Linear Layer :',fc_2)Â
# APPLY THE Second Linear layer to the output y3y4 = fc_2(y3)print('Output Shape :',y4.shape)Â
# Computational Graphmake_dot(y4, params=dict(fc_2.named_parameters())) |
Output:
Second Linear Layer : Linear(in_features=120, out_features=1, bias=True) Output Shape : torch.Size([1, 16, 358, 1])
Computational Graph
Step 8: Train the model and plot the final computational graph
Python3
# define labelslabels = torch.rand(1).to(device)Â
# forward propagationprediction = net(img)Â
# define lossloss = (prediction - labels).sum()Â
# backward passloss.backward() # Optimizeroptimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)Â
# gradient descentoptimizer.step() Â
# Computational Graphmake_dot(prediction, Â Â Â Â Â Â Â Â Â params=dict(net.named_parameters())) |
Output:
Computational Graph
