Model Graphs

Supported Frameworks

We currently support PyTorch and Lightning.

mlop supports logging model details to visualize the model architecture, including layers, connections, parameters, and more.

Logging Model Details

Currently, mlop supports directly invoking mlop.watch() on a model to log the model graph.

mlop.watch(
    model,
    disable_graph=False,
    **kwargs,
)

Data Flow

In PyTorch/Lightning, model details are recorded by tracing the autograd function across every module in the model. This is done once upon initialization. The first iteration of the model will be used to record the data flow and tree structure of the model.

Disabling Data Flow Logging

In certain models, the autograd functions injected may affect the ability to save onnx models. In this case, you can disable data flow logging by settingdisable_graph=True.

Example

import torch
import torch.nn as nn
import mlop
 
class ConvNet(nn.Module):
    def __init__(self, kernels, classes=10):
        super(ConvNet, self).__init__()
 
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, kernels[0], kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, kernels[1], kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.fc = nn.Linear(7 * 7 * kernels[-1], classes)
 
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out
 
model = ConvNet([16, 32], 10)
mlop.watch(model, disable_graph=False, freq=100)

This provides you with nice visualization as soon as the model training starts,

Visualizing PyTorch/Lightning Models

For each node, mlop can log the following details for both the input and the output:

Name
Type
Position
Shape
Parameters
Connections

Model Graphs

Logging Model Details

Example

Visualizing PyTorch/Lightning Models

On this page