Skip to main content
Product
Developers
Connect
Company
Log inGet started

Visualize a model

The max visualize command helps you inspect your model after it’s been parsed and partially lowered into the “MO” dialect (an intermediate representation for graph ops that's used by the MAX Engine compiler). The command outputs a .maxviz file that you can open in Netron to inspect the graph.

Visualizing your model as a “MO” graph is useful when you want to explore the graph layout after early decomposition in the MAX Engine compiler. It can also help you identify ops that fallback to the stock framework implementations—ops that might be suitable for you to replace with a custom op (coming soon).

Figure 1. Screenshot of a maxviz file in Netron.

caution

The graph in a maxviz file is a deconstructed version of the model before optimization in the MAX Engine compiler, so it might be much larger than the original graph. That means Netron will struggle to load models that are very complex, and it might crash. The Netron community is aware of this general issue, but we are also looking for opportunities to simplify the maxviz file to avoid this as often as possible.

Create a maxviz file​

Pass any TorchScript or ONNX file to max visualize and it will generate a .maxviz file:

max visualize path/to/model

By default, this saves the .maxviz file to the local directory, but you can specify another directory and the filename with --output.

Then, just drag and drop the .maxviz file into Netron.

The visualize command actually compiles the model (partially), so it needs to know the model's input shapes. These shapes specified by metadata in a TF SavedModel or ONNX file. However, TorchScript models don't provide such metadata. If you're using a TorchScript model, then you need to specify the input shape, as we'll show next.

note

The .maxviz file is not a fully-compiled model and it is not an executable graph. It represents the high-level graph layout in the “MO” dialect and it’s designed only for visualization in Netron.

Specify the input shapes​

If your model doesn’t provide input metadata, you must specify the input shapes with a YAML file that you pass to the --input-data-schema option.

For example, let’s say you have ResNet-50 as a TorchScript model. It has an input shape of 3x224x224, plus a batch size, and it uses float-32 data. For the purposes of our tool, you need to pick a fixed batch size, so let’s use 1. That translates to an input shape of 1x3x224x224xf32. Finally, you need to specify the input tensor name (in this example, it's pixel_values). So, your YAML file looks like this:

input-spec.yaml
inputs:
  - input_name: pixel_values
    shape: 1x3x224x224xf32

For demonstration, we've prepared this for you in our GitHub repo and you can try it like this:

  1. Download the ResNet-50 PyTorch model and convert it to TorchScript:

    cd max/examples/tools/common/resnet50-pytorch/
    bash download-model.sh --output resnet50.torchscript

  2. Then call the visualize command and pass the input-spec.yaml file that's already there:

    max visualize resnet50.torchscript --input-data-schema input-spec.yaml

    This creates a .maxviz file in the local directory.

  3. Now, drag and drop the .maxviz file into Netron.

For more information about this input data schema, see the input schema reference.

What’s in the maxviz file​

The maxviz file is a deconstructed version of the model before it’s optimized in the MAX Engine compiler. So, it is not a fully-compiled model and it is not executable. The maxviz file represents the high-level graph layout in the “MO” dialect and it’s designed purely for visualization in Netron.

When you click on a node in Netron, you can see whether that op is implemented natively as a “MO” op. If it is, the is_fallback property is False and the node type/name begins with mo.

If is_fallback is True, then the graph includes additional tensor-translation nodes before and after the op, which convert the tensor between the “MO” tensor format and the stock framework tensor format. For example, figure 2 shows an adaptve_avg_pool2d op that is a fallback. The input to the fallback is translated from "mo" dialect into the fallback op's dialect using an node named m_tensor_to_ivalue_list. Likewise, the fallback's output is translated back into "mo" dialect.

Figure 2. Close-up of a graph that includes a fallback op.

These fallback ops ensure that your model is still compatible with MAX Engine, and you can still receive significant performance gains when not all model ops are available in the MO dialect. If you want to optimize the model further, then you provide your own native implementation of the op written in Mojo (coming soon).

Was this page helpful?