Export from PyTorch
In general, the procedure for model export is pretty straightforward thanks to good integration of
.onnx in PyTorch.
The code itself is simple. First we import torch and build a test model.
import torch import torch.nn as nn class Model(nn.Module): def __init__(self): super().__init__() self.layer = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, stride=2, padding=1) def forward(self, x): out = self.layer(x) out = nn.functional.interpolate(out, scale_factor=2, mode='bilinear', align_corners=True) out = torch.nn.functional.softmax(out, dim=1) return out
The code for the export itself is then
model = Model() model.eval() random_input = torch.randn(1, 3, 64, 64, dtype=torch.float32) # you can add however many inputs your model or task requires input_names = ["image"] output_names = ["pred"] torch.onnx.export(model, random_input, './model.onnx', verbose=False, input_names=input_names, output_names=output_names, opset_version=11)
It is important to make sure that the number of elements in
input_names is the same as the number of input arguments in your model’s
forward method. As well as that the number of return variables of the forward method is the same as the number of elements in
Make sure that input shapes are correct since you won’t be able to change them when you import the model in Lens Studio. It just takes them from your
Tip: You can download this guide as .ipynb here
Loading Models into Lens Studio
The process for importing your model into LensStudio is again straightforward. You just need to add an ML component and it will prompt you to select a file containing your model. Select your
ONNX file that you’ve exported previously and if everything is fine, the studio will prompt you to set your model’s input and output scale and bias.
In the upper left corner of this prompt you will see a compatibility table of Ops in your model and different available inference frameworks. These are used to accelerate model inference if the device that’s running your lens has the ability to utilize these frameworks and if all of the Ops are implemented in a given inference framework. If your model has unsupported Ops, you can hover over the warning icon and it will provide you with a layer name that has this issue.
Tip: Take a look at the Compatibility guide to learn more about supported Layers
One of the most useful apps for debugging ONNX files is Netron.
You can open your
.onnx file in Netron and find the Op that’s causing the issue by name.
You can then see if you can reimplement your model avoiding unsupported operations.
Remember that Lens Studio textures have values in the range of
[0, 255] so if your mode was trained with whitened input with values in
[-1, 1] you need to correctly set scale and bias parameters to reflect that.
Tip: Check out the ML Component guide for more information!
One of the major issues is bilinear interpolation. There is a discrepancy between PyTorch and mobile inference frameworks in handling edges of interpolated image with
align_corners set to
False. So you need to make sure your model uses
align_corners=True everywhere it uses bilinear interpolation. Also you should use
opset_version=11 if you have
align_corners=True in your model, since default 9th opset doesn’t have this parameter in its op definition. It might also be better to use nearest neighbor interpolation or transposed convolution instead of bilinear interpolation, since iOS inference accelerator doesn’t support
x = random_input # bad nn.functional.interpolate(x, scale_factor=2, mode='bilinear') nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False) nn.Upsample(scale_factor=2, mode='bilinear') # better nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True) nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) # best nn.ConvTranspose2d(3, 3, 3) nn.functional.interpolate(x, scale_factor=2, mode='nearest') nn.Upsample(scale_factor=2, mode='nearest')
Pay attention to your
.reshape() operations in your network. Main issue with those is that while PyTorch uses NCHW format for its 4 dimensional tensors, different mobile inference frameworks can use either NCHW or NHWC format. If possible replace them with operations that preserve tensor dimensions.
# bad: NCHW dim order is assumed features = features.view(batch_size, num_channels, height, width) result = conv(features) # ok: dim order is irrelevant flattened_features = features.view(batch_size, -1) result = flattened_features.sum(dim=1)
Some PyTorch versions have issues with different ONNX opsets. If you are encountering issues exporting model with interpolation, softmax layer with set dim parameter, try to update your PyTorch to the latest available version and set
opset_version=11 parameter in your
torch.onnx.export function call.
Still Looking for help?Visit Support