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.
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.
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:
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:
-
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
-
Then call the
visualize
command and pass theinput-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. -
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?
Thank you! We'll create more content like this.
Thank you for helping us improve!
If you'd like to share more information, please report an issue on GitHub
😔 What went wrong?