Unable to compile PyTorch quantized models with torch-mlir due to CustomClass #4286
Description
What happened?
I'm trying to convert a quantized ResNet-50 model (using PyTorch's quantization-aware training) to MLIR format, but I'm encountering an error when torch-mlir attempts to import the model.
Environment
- python:
3.11.3 - pytorch:
2.7.1+cpu - torch-mlir:
1983b6db9cdba2dfef4939a4fb521e9ac25dc08e
Steps to reproduce the issue
import torch
import torch.nn as nn
import torch.quantization
import torchvision
import torchvision.transforms as transforms
import torch_mlir.torchscript
import os
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_dataset = torchvision.datasets.FakeData(transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
model = torchvision.models.quantization.resnet50(pretrained=True, quantize=False)
model.qconfig = torch.quantization.get_default_qat_qconfig("fbgemm")
torch.quantization.prepare_qat(model, inplace=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.train()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
# QAT
for epoch in range(1):
for images, targets in train_loader:
images, targets = images.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
print(f"[Epoch {epoch+1}] QAT training complete")
model.eval()
quantized_model = torch.quantization.convert(model.eval(), inplace=False)
print(quantized_model)
scripted_model = torch.jit.script(quantized_model)
example_input = torch.randn(1, 3, 224, 224)
mlir_module = torch_mlir.torchscript.compile(
scripted_model,
[example_input],
output_type="torch"
)
with open('resnet50_int8.mlir', 'w') as f:
f.write(str(mlir_module))log:
QuantizableResNet(
(conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU()
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): QuantizableBottleneck(
(conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU()
(downsample): Sequential(
(0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(skip_add_relu): FloatFunctional(
(activation_post_process): Identity()
)
(relu1): ReLU()
(relu2): ReLU()
)
...
Traceback (most recent call last):
File "/workspace/qat_resnet50.py", line 57, in <module>
mlir_module = torch_mlir.torchscript.compile(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/path/to/torch-mlir/build/python_packages/torch_mlir/torch_mlir/torchscript.py", line 282, in compile
raise Exception(
Exception:
PyTorch TorchScript module -> torch-mlir Object Graph IR import failed with:
### Importer C++ Exception:
see diagnostics
### Importer Diagnostics:
error: unable to import Torch CustomClass type '0x5baca5da1580' to MLIR type
Do we have any approach to convert a quantized model to MLIR currently?