ERROR - converting 'full' op (located at: 'heads')

[samxu29]

🐞Describing the bug

Get KeyError: 6 when converting Yolo Nas Pose model

[2024-07-07 18:31:13] ERROR - ops.py - 

ERROR - converting 'full' op (located at: 'heads'):

Converting PyTorch Frontend ==> MIL Ops:  94%|██████████████████████████████████████████████████████████▏   | 1142/1217 [00:00<00:00, 3450.93 ops/s]
Traceback (most recent call last):
  File "/Users/{path}/yolo_nas/export_superGradients.py", line 25, in <module>
    coreml_model = ct.convert(
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/_converters_entry.py", line 581, in convert
    mlmodel = mil_convert(
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 188, in mil_convert
    return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 212, in _mil_convert
    proto, mil_program = mil_convert_to_proto(
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 288, in mil_convert_to_proto
    prog = frontend_converter(model, **kwargs)
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 108, in __call__
    return load(*args, **kwargs)
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 82, in load
    return _perform_torch_convert(converter, debug)
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 116, in _perform_torch_convert
    prog = converter.convert()
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 581, in convert
    convert_nodes(self.context, self.graph)
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 86, in convert_nodes
    raise e     # re-raise exception
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 81, in convert_nodes
    convert_single_node(context, node)
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 134, in convert_single_node
    add_op(context, node)
  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 4220, in full
    else NUM_TO_NUMPY_DTYPE[TORCH_DTYPE_TO_NUM[inputs[2].val]]
KeyError: 6

To Reproduce

import torch
from super_gradients.common.object_names import Models
from super_gradients.conversion import DetectionOutputFormatMode
from super_gradients.training import models
import coremltools as ct

model_type = 'yolo_nas_pose_n'

model = models.get(Models.YOLO_NAS_POSE_N, pretrained_weights="coco_pose")
export_result = model.export(f"yolo_nas/models_onnx/{model_type}.onnx",
                             postprocessing=True,
                             output_predictions_format=DetectionOutputFormatMode.BATCH_FORMAT
                             )
# coreml conversion
model.prep_model_for_conversion(input_size=[1, 3, 640, 640])
torch.save(model, "yolo_nas/models_torch/yolo_nas_pose_n.pth")

traced_model = torch.jit.trace(model, torch.randn(1, 3, 640, 640))

torch.jit.save(traced_model, f"yolo_nas/models_torch/{model_type}_traced.pth")

input_image = ct.ImageType(name="my_input", shape=(1, 3, 640, 640), scale=1/255)

# Convert the models coreml
coreml_model = ct.convert(
    traced_model,
    convert_to='mlprogram',
    inputs=[input_image],
    minimum_deployment_target=ct.target.iOS17,
    debug=True  # Enable debugging
)

coreml_model.save(f"yolo_nas/models_coreml/{model_type}.mlpackage")

System environment (please complete the following information):

• coremltools version: 7.2
• OS (e.g. MacOS version or Linux type): MacOS 14.5
• Any other relevant version information (e.g. PyTorch or TensorFlow version): torch 2.2.0

Additional context

• Add anything else about the problem here that you want to share.

[YifanShenSZ]

Hey @samxu29, could you please try our latest coremltools 8.0b1?

Concretely, the error line

  File "/Users/{path}/.venv/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 4220, in full
    else NUM_TO_NUMPY_DTYPE[TORCH_DTYPE_TO_NUM[inputs[2].val]]

no longer exists in tot source code, so presumably it has been fixed

[samxu29]

The issues is solved in 8.0b1 version, thank you for your help!