Does timm ViT models support PyTorch to ONNX conversion?

[manideep2510]

Hello

Can we convert PyTorch ViT models provided in timm to ONNX.

When I tried to convert vit_base_patch16_224 model to ONNX, I am getting the below warning and the output of PyTorch and ONNX models are different for the same input.

TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert H == self.img_size[0] and W == self.img_size[1], \

Thanks in advance.

[alexander-soare]

Don't know much about ONNX myself, but in case it's useful, I know that tracer warning isn't your culprit. If you use an input with a valid size, that assert should never need to trigger. This warning would only matter if there were control flow with logic that has an impact on the output.

[alexander-soare]

I was able to get a good result with (following this tutorial):

import onnxruntime

import torch
import timm
import numpy as np

model = timm.create_model('vit_base_patch16_224', pretrained=True)
model.eval()

inp = torch.randn(2, 3, 224, 224)
with torch.no_grad():
    torch_out = model(inp)

torch.onnx.export(model, inp, 'test.onnx', export_params=True, input_names = ['input'], output_names = ['output'])

ort_session = onnxruntime.InferenceSession('test.onnx')

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# Compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(inp)}
ort_outs = ort_session.run(None, ort_inputs)

# Compare ONNX Runtime and PyTorch results
if np.allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05):
    print("Exported model has been tested with ONNXRuntime, and the result looks good!")
else:
    print("FAILED due to mismatch between outputs")

[manideep2510]

Thanks a lot for the help @alexander-soare

I was doing a blunder while testing the inference. I was using different num_classes while initializing the model with imagenet weights and converting that to ONNX. But while testing the inference and model outputs, I was just initializing the timm PyTorch model with imagenet weights and testing it against the ONNX model. But this is not correct, because, every time I initialize the model with num_classes other than it's default value, the last layer's weights will be randomly initialized every time and that model is no longer the same as the model used for ONNX conversion.

So when I save the pytorch weights and loaded them while testing, the onnx and pytorch outputs match perfectly.

model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=10)
model.load_state_dict(torch.load('test.pth'))

[rwightman]

FYI, you can just
model = timm.create_model('vit_base_patch16_224', num_classes=10, checkpoint_path='test.pth')