I've trained a YOLO-NAS model on Roboflow , nut unable to get it quantized and conveted to ONNX using the MCT library. I am newbie and all the code I could come up with gives me the following error
I am using the following code -
import torch
from torchvision import transforms
from PIL import Image
import os
import model_compression_toolkit as mct
#from model_compression_toolkit.target_platform_capabilities import get_target_platform_capabilities
# --- Config ---
CKPT_PATH = "test.pt" # your checkpoint file
IMG_SIZE = 416
REPRESENTATIVE_DATA_DIR = "dataset/images/val"
ONNX_EXPORT_PATH = "yolo_nas_s_quantized.onnx"
# --- Load YOLO NAS S model from SuperGradients ---
from super_gradients.training import models
# You must specify num_classes or pretrained_weights according to your model
NUM_CLASSES = 1 # change if your dataset is different
model = models.get("yolo_nas_s", num_classes=NUM_CLASSES, checkpoint_path=CKPT_PATH)
model.eval()
# Prepare the model for deployment. This converts dynamic operations (like anchor generation)
# into a static, traceable format required for quantization and export.
model.prep_model_for_conversion(input_size=(IMG_SIZE, IMG_SIZE))
# --- Representative data generator ---
preprocess = transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ToTensor(),
])
# Filter for common image file extensions and get the count for calibration
img_files = [f for f in os.listdir(REPRESENTATIVE_DATA_DIR) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
def representative_data_gen():
for img_name in img_files:
img_path = os.path.join(REPRESENTATIVE_DATA_DIR, img_name)
img = Image.open(img_path).convert("RGB")
input_tensor = preprocess(img).unsqueeze(0) # add batch dim
yield [input_tensor]
# --- Get target platform capabilities for PyTorch default ---
#tpc = get_target_platform_capabilities('pytorch', 'default')
tpc = mct.get_target_platform_capabilities('pytorch', 'default')
# With a clean environment, you'll have the latest MCT version.
# It expects n_calibration_iter as a direct argument.
core_config = mct.core.CoreConfig()
# --- Run post training quantization ---
quantized_model, quant_info = mct.ptq.pytorch_post_training_quantization(
in_module=model,
representative_data_gen=representative_data_gen,
target_platform_capabilities=tpc
)
print("Quantization complete.")
# --- Export quantized model to ONNX ---
# Use the MCT exporter which is aware of the quantization scheme.
# It correctly translates the quantized model to an ONNX graph.
mct.exporter.pytorch_export_model(model=quantized_model, save_model_path=ONNX_EXPORT_PATH, repr_dataset=representative_data_gen, opset_version=13)
print(f"Quantized ONNX model exported to: {ONNX_EXPORT_PATH}")
And getting the following error
File "/Users/ashishol/My_Code/Python/Yolo/quantize_yolo1.py", line 51, in <module>
quantized_model, quant_info = mct.ptq.pytorch_post_training_quantization(
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/ptq/pytorch/quantization_facade.py", line 125, in pytorch_post_training_quantization
tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_module,
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/runner.py", line 99, in core_runner
graph = graph_preparation_runner(in_model,
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/graph_prep_runner.py", line 73, in graph_preparation_runner
graph = read_model_to_graph(in_model,
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/graph_prep_runner.py", line 211, in read_model_to_graph
graph = fw_impl.model_reader(in_model,
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/pytorch_implementation.py", line 175, in model_reader
return model_reader(_module, representative_data_gen, self.to_numpy, self.to_tensor)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/reader/reader.py", line 174, in model_reader
fx_model = fx_graph_module_generation(model, representative_data_gen, to_tensor)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/reader/reader.py", line 111, in fx_graph_module_generation
symbolic_traced = _trace_model(pytorch_model)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/reader/reader.py", line 46, in _trace_model
graph = tracer.trace(root)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 837, in trace
(self.create_arg(fn(*args)),),
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/super_gradients/training/models/detection_models/customizable_detector.py", line 95, in forward
return self.heads(x)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 812, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 530, in call_module
ret_val = forward(*args, **kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 805, in forward
return _orig_module_call(mod, *args, **kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/super_gradients/training/models/detection_models/yolo_nas/dfl_heads.py", line 227, in forward
anchor_points_inference, stride_tensor = self._generate_anchors(feats)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/super_gradients/training/models/detection_models/yolo_nas/dfl_heads.py", line 263, in _generate_anchors
shift_x = torch.arange(end=w, dtype=torch.float32, device=device) + self.grid_cell_offset
TypeError: arange() received an invalid combination of arguments - got (device=Attribute, dtype=torch.dtype, end=Proxy, ), but expected one of:
* (Number end, *, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
* (Number start, Number end, *, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
* (Number start, Number end, Number step = 1, *, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
Will appreciate any help to proceed.
Thanks in advance,
I am using the following code -
import torch
from torchvision import transforms
from PIL import Image
import os
import model_compression_toolkit as mct
#from model_compression_toolkit.target_platform_capabilities import get_target_platform_capabilities
# --- Config ---
CKPT_PATH = "test.pt" # your checkpoint file
IMG_SIZE = 416
REPRESENTATIVE_DATA_DIR = "dataset/images/val"
ONNX_EXPORT_PATH = "yolo_nas_s_quantized.onnx"
# --- Load YOLO NAS S model from SuperGradients ---
from super_gradients.training import models
# You must specify num_classes or pretrained_weights according to your model
NUM_CLASSES = 1 # change if your dataset is different
model = models.get("yolo_nas_s", num_classes=NUM_CLASSES, checkpoint_path=CKPT_PATH)
model.eval()
# Prepare the model for deployment. This converts dynamic operations (like anchor generation)
# into a static, traceable format required for quantization and export.
model.prep_model_for_conversion(input_size=(IMG_SIZE, IMG_SIZE))
# --- Representative data generator ---
preprocess = transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ToTensor(),
])
# Filter for common image file extensions and get the count for calibration
img_files = [f for f in os.listdir(REPRESENTATIVE_DATA_DIR) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
def representative_data_gen():
for img_name in img_files:
img_path = os.path.join(REPRESENTATIVE_DATA_DIR, img_name)
img = Image.open(img_path).convert("RGB")
input_tensor = preprocess(img).unsqueeze(0) # add batch dim
yield [input_tensor]
# --- Get target platform capabilities for PyTorch default ---
#tpc = get_target_platform_capabilities('pytorch', 'default')
tpc = mct.get_target_platform_capabilities('pytorch', 'default')
# With a clean environment, you'll have the latest MCT version.
# It expects n_calibration_iter as a direct argument.
core_config = mct.core.CoreConfig()
# --- Run post training quantization ---
quantized_model, quant_info = mct.ptq.pytorch_post_training_quantization(
in_module=model,
representative_data_gen=representative_data_gen,
target_platform_capabilities=tpc
)
print("Quantization complete.")
# --- Export quantized model to ONNX ---
# Use the MCT exporter which is aware of the quantization scheme.
# It correctly translates the quantized model to an ONNX graph.
mct.exporter.pytorch_export_model(model=quantized_model, save_model_path=ONNX_EXPORT_PATH, repr_dataset=representative_data_gen, opset_version=13)
print(f"Quantized ONNX model exported to: {ONNX_EXPORT_PATH}")
And getting the following error
File "/Users/ashishol/My_Code/Python/Yolo/quantize_yolo1.py", line 51, in <module>
quantized_model, quant_info = mct.ptq.pytorch_post_training_quantization(
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/ptq/pytorch/quantization_facade.py", line 125, in pytorch_post_training_quantization
tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_module,
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/runner.py", line 99, in core_runner
graph = graph_preparation_runner(in_model,
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/graph_prep_runner.py", line 73, in graph_preparation_runner
graph = read_model_to_graph(in_model,
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/graph_prep_runner.py", line 211, in read_model_to_graph
graph = fw_impl.model_reader(in_model,
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/pytorch_implementation.py", line 175, in model_reader
return model_reader(_module, representative_data_gen, self.to_numpy, self.to_tensor)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/reader/reader.py", line 174, in model_reader
fx_model = fx_graph_module_generation(model, representative_data_gen, to_tensor)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/reader/reader.py", line 111, in fx_graph_module_generation
symbolic_traced = _trace_model(pytorch_model)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/reader/reader.py", line 46, in _trace_model
graph = tracer.trace(root)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 837, in trace
(self.create_arg(fn(*args)),),
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/super_gradients/training/models/detection_models/customizable_detector.py", line 95, in forward
return self.heads(x)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 812, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 530, in call_module
ret_val = forward(*args, **kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 805, in forward
return _orig_module_call(mod, *args, **kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/super_gradients/training/models/detection_models/yolo_nas/dfl_heads.py", line 227, in forward
anchor_points_inference, stride_tensor = self._generate_anchors(feats)
File "/opt/homebrew/Caskroom/miniforge/base/envs/yolo-quant/lib/python3.10/site-packages/super_gradients/training/models/detection_models/yolo_nas/dfl_heads.py", line 263, in _generate_anchors
shift_x = torch.arange(end=w, dtype=torch.float32, device=device) + self.grid_cell_offset
TypeError: arange() received an invalid combination of arguments - got (device=Attribute, dtype=torch.dtype, end=Proxy, ), but expected one of:
* (Number end, *, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
* (Number start, Number end, *, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
* (Number start, Number end, Number step = 1, *, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
Will appreciate any help to proceed.
Thanks in advance,
Statistics: Posted by ashishol — Tue Jul 01, 2025 3:57 am — Replies 1 — Views 30