ML model visualization using Netron

1. Intro

• Netron is really great tool to visualize the ML models.
• ONNX (Open Neural Network Exchange Format) is a format designed to represent any type of Machine Learning and Deep Learning models.
• Pytorch and Tensorflow can export the model in onnx format.
• This onnx file can then be read with Netron to visualize the model graphically.

2. Installation

As simple as:

pip install netron

3. Barebone code

import torch
import torch.nn as nn
import netron

# Define the Model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.dense1 = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
            nn.ReLU()
        )    

    def forward(self, x):
        x = self.flatten(x)
        out = self.dense1(x)
        return out

model = NeuralNetwork() # Create a Model
dummy_input = torch.ones([1,28,28]) # Dummy Input to the model 
									# (batch_size=1, input_shape=(28x28))
model_path = "simple_model.onnx" # Path of the onnx file
torch.onnx.export(model, dummy_input, model_path) # Export the model to ONNX file
netron.start(model_path) # Visualize using Netron

Output

• Above code will open the visualization in browser @ http://localhost:8080/
• It needs a dummy input to pass through the model and calculate necessary data.
• Dummy input can be anything. It just needs to have right shape that model takes.
• To name the inputs and outputs, add following parameters to the torch.onnx.export function.

torch.onnx.export(model, dummy_input, model_path,
                  input_names = ['image'],   
                  output_names = ['prediction'])

4. Show tensor shapes at each step

• Install onnx package

  pip install onnx

• After exporting the onnx file, read back using onnx.load and then infer the tensor shapes

  # Input tensor shapes at each step
  onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model_path)), model_path)
  
  netron.start(model_path) # Visualize using Netron

Output

• This is a lot better visualization of model

5. Multiple Inputs Outputs

• Here is example using Two inputs and Two outputs

  import torch
  import torch.nn as nn
  import netron
  import onnx
  
  # Define the Model
  class NeuralNetwork(nn.Module):
      def __init__(self):
          super(NeuralNetwork, self).__init__()
          self.flatten = nn.Flatten()
          self.dense1 = nn.Sequential(
              nn.Linear(28*28, 512),
              nn.ReLU(),
              nn.Linear(512, 10),
              nn.ReLU()
          ) 
          self.dense2 = nn.Sequential(
              nn.Linear(14*14, 256),
              nn.ReLU(),
              nn.Linear(256, 10),
              nn.ReLU()
          ) 
          self.head1 = nn.Sequential(
              nn.Linear(20,10),
              nn.ReLU()
          ) 
          self.head2 = nn.Sequential(
              nn.Linear(20,15),
              nn.ReLU(),
              nn.Linear(15,10),
              nn.ReLU()
          )  
  
      def forward(self, x, y):
          x_out = self.flatten(x)
          y_out = self.flatten(y)
  
          x_out = self.dense1(x_out)
          y_out = self.dense2(y_out)
  
          combined = torch.cat((x_out,y_out),1)
  
          pred1 = self.head1(combined)
          pred2 = self.head2(combined)
  
          return pred1, pred2
  
  model = NeuralNetwork() # Create a Model
  dummy_input1 = torch.ones([1,28,28]) # Dummy Input (batch_size=1, input_shape=(28x28))
  dummy_input2 = torch.ones([1,14,14]) # Dummy Input (batch_size=1, input_shape=(14x14))
  
  model_path = "simple_model.onnx" # Path of the onnx file
  torch.onnx.export(model, (dummy_input1, dummy_input2), model_path,
                    input_names = ['image1', 'image2'],   
                    output_names = ['prediction1', 'prediction2'])  # Export the model to ONNX file
  
  # Input tensor shapes at each step
  onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model_path)), model_path)
  
  netron.start(model_path) # Visualize using Netron

Output