Error converting from torch jit trace

[darwinharianto]

import torch
import numpy as np
from transformers import AutoModel, AutoProcessor, OwlViTModel, OwlViTProcessor, CLIPTokenizerFast, CLIPTokenizer
import coremltools as ct
from PIL import Image
import requests

class MyOpenDetector(torch.nn.Module):
    def __init__(self, model=None):
        super(MyOpenDetector, self).__init__()
        self.model = model

    def forward(self, input_ids, pixel_values, attention_mask):
        # inputs = {"input_ids":x[0], "attention_mask":x[1], "pixel_values":x[2]}
        outputs = self.model(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask)
        logits_per_image = outputs[0]  # this is the image-text similarity score
        probs = logits_per_image.softmax(dim=1)  # we can take the softmax to get the label probabilities

        return probs

def save_owlvitmodel(inputs, modelname):

    openModel = AutoModel.from_pretrained(modelname, torchscript=True).eval()

    x = tuple([inputs['input_ids'], inputs['pixel_values'], inputs['attention_mask']])
    model = MyOpenDetector(model=openModel)
    traced_model = torch.jit.trace(model, x)
    torch.jit.save(traced_model, 'traced_owlvit.pt')

    return traced_model

modelname = "google/owlvit-base-patch32"
processor = AutoProcessor.from_pretrained(modelname, torchscript=True)
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(text=[["a photo of a cat", "a photo of a dog"]], images=torch.Tensor(np.asarray(image)), return_tensors="pt")

traced_model = save_owlvitmodel(inputs, modelname)

loaded_model = torch.jit.load("traced_owlvit.pt")
loaded_model.eval()

x = tuple([inputs['input_ids'], inputs['pixel_values'], inputs['attention_mask']])
probs = loaded_model(*x)

mlmodel = ct.convert(
    traced_model,
    inputs=[ct.TensorType(name="input_ids", shape=(ct.RangeDim(1, 16),16), dtype=np.int32),
            ct.TensorType(name="pixel_values", shape=(ct.RangeDim(1, 3),3,768,768), dtype=np.float32),
            ct.TensorType(name="attention_mask", shape=(ct.RangeDim(1, 16),16), dtype=np.int32),],
)
mlmodel.save('coremlmodel_owlvit.mlmodel')

System environment (please complete the following information):

• coremltools version:
• Ubuntu 20
• torch==2.0.0

[TobyRoseman]

Stack Trace

``` ---> 1 mlmodel = ct.convert( 2 traced_model, 3 inputs=[ct.TensorType(name="input_ids", shape=(ct.RangeDim(1, 16),16), dtype=np.int32), 4 ct.TensorType(name="pixel_values", shape=(ct.RangeDim(1, 3),3,768,768), dtype=np.float32), 5 ct.TensorType(name="attention_mask", shape=(ct.RangeDim(1, 16),16), dtype=np.int32),], 6 ) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/_converters_entry.py:492, in convert(model, source, inputs, outputs, classifier_config, minimum_deployment_target, convert_to, compute_precision, skip_model_load, compute_units, package_dir, debug, pass_pipeline) 489 if specification_version is None: 490 specification_version = _set_default_specification_version(exact_target) --> 492 mlmodel = mil_convert( 493 model, 494 convert_from=exact_source, 495 convert_to=exact_target, 496 inputs=inputs, 497 outputs=outputs_as_tensor_or_image_types, # None or list[ct.ImageType/ct.TensorType] 498 classifier_config=classifier_config, 499 skip_model_load=skip_model_load, 500 compute_units=compute_units, 501 package_dir=package_dir, 502 debug=debug, 503 specification_version=specification_version, 504 main_pipeline=pass_pipeline, 505 ) 507 if exact_target == 'milinternal': 508 return mlmodel # Returns the MIL program File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/converter.py:188, in mil_convert(model, convert_from, convert_to, compute_units, **kwargs) 149 @_profile 150 def mil_convert( 151 model, (...) 155 **kwargs 156 ): 157 """ 158 Convert model from a specified frontend `convert_from` to a specified 159 converter backend `convert_to`. (...) 186 See `coremltools.converters.convert` 187 """ --> 188 return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/converter.py:212, in _mil_convert(model, convert_from, convert_to, registry, modelClass, compute_units, **kwargs) 209 weights_dir = _tempfile.TemporaryDirectory() 210 kwargs["weights_dir"] = weights_dir.name --> 212 proto, mil_program = mil_convert_to_proto( 213 model, 214 convert_from, 215 convert_to, 216 registry, 217 **kwargs 218 ) 220 _reset_conversion_state() 222 if convert_to == 'milinternal': File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/converter.py:285, in mil_convert_to_proto(model, convert_from, convert_to, converter_registry, main_pipeline, **kwargs) 280 frontend_pipeline, backend_pipeline = _construct_other_pipelines( 281 main_pipeline, convert_from, convert_to 282 ) 284 frontend_converter = frontend_converter_type() --> 285 prog = frontend_converter(model, **kwargs) 286 PipelineManager.apply_pipeline(prog, frontend_pipeline) 288 PipelineManager.apply_pipeline(prog, main_pipeline) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/converter.py:108, in TorchFrontend.__call__(self, *args, **kwargs) 105 def __call__(self, *args, **kwargs): 106 from .frontend.torch.load import load --> 108 return load(*args, **kwargs) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py:63, in load(model_spec, inputs, specification_version, debug, outputs, cut_at_symbols, **kwargs) 55 inputs = _convert_to_torch_inputtype(inputs) 56 converter = TorchConverter( 57 torchscript, 58 inputs, (...) 61 specification_version, 62 ) ---> 63 return _perform_torch_convert(converter, debug) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py:102, in _perform_torch_convert(converter, debug) 100 def _perform_torch_convert(converter, debug): 101 try: --> 102 prog = converter.convert() 103 except RuntimeError as e: 104 if debug and "convert function" in str(e): File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/converter.py:284, in TorchConverter.convert(self) 281 self.convert_const() 283 # Add the rest of the operations --> 284 convert_nodes(self.context, self.graph) 286 graph_outputs = [self.context[name] for name in self.graph.outputs] 288 # An output can be None when it's a None constant, which happens 289 # in Fairseq MT. File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py:88, in convert_nodes(context, graph) 83 raise RuntimeError( 84 "PyTorch convert function for op '{}' not implemented.".format(node.kind) 85 ) 87 context.prepare_for_conversion(node) ---> 88 add_op(context, node) 90 # We've generated all the outputs the graph needs, terminate conversion. 91 if _all_outputs_present(context, graph): File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py:3315, in _internal_op_tensor_inplace_fill(context, node) 3312 data = context[node.inputs[0]] 3313 fill_scalar = context[node.inputs[1]] -> 3315 begin, end, stride, begin_mask, end_mask, squeeze_mask = _get_slice_params( 3316 context, data, node.inputs[2:] 3317 ) 3318 if begin.val is None or end.val is None: 3319 raise ValueError("_internal_op_tensor_inplace_fill does not support dynamic index") File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py:3240, in _get_slice_params(context, data, inputs) 3239 def _get_slice_params(context, data, inputs): -> 3240 rank = data.rank 3241 begin = [0] * rank 3242 end = [0] * rank AttributeError: 'list' object has no attribute 'rank' ```

It also fails without using flexible shape (i.e. not using ct.RangeDim, just using the shapes in x). Although the error is different.

Stack Trace

``` ----> 1 mlmodel = ct.convert( 2 traced_model, 3 inputs=[ct.TensorType(name="input_ids", shape=(2,16), dtype=np.int32), 4 ct.TensorType(name="pixel_values", shape=(1,3,768,768), dtype=np.float32), 5 ct.TensorType(name="attention_mask", shape=(2,16), dtype=np.int32),], 6 ) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/_converters_entry.py:492, in convert(model, source, inputs, outputs, classifier_config, minimum_deployment_target, convert_to, compute_precision, skip_model_load, compute_units, package_dir, debug, pass_pipeline) 489 if specification_version is None: 490 specification_version = _set_default_specification_version(exact_target) --> 492 mlmodel = mil_convert( 493 model, 494 convert_from=exact_source, 495 convert_to=exact_target, 496 inputs=inputs, 497 outputs=outputs_as_tensor_or_image_types, # None or list[ct.ImageType/ct.TensorType] 498 classifier_config=classifier_config, 499 skip_model_load=skip_model_load, 500 compute_units=compute_units, 501 package_dir=package_dir, 502 debug=debug, 503 specification_version=specification_version, 504 main_pipeline=pass_pipeline, 505 ) 507 if exact_target == 'milinternal': 508 return mlmodel # Returns the MIL program File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/converter.py:188, in mil_convert(model, convert_from, convert_to, compute_units, **kwargs) 149 @_profile 150 def mil_convert( 151 model, (...) 155 **kwargs 156 ): 157 """ 158 Convert model from a specified frontend `convert_from` to a specified 159 converter backend `convert_to`. (...) 186 See `coremltools.converters.convert` 187 """ --> 188 return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/converter.py:212, in _mil_convert(model, convert_from, convert_to, registry, modelClass, compute_units, **kwargs) 209 weights_dir = _tempfile.TemporaryDirectory() 210 kwargs["weights_dir"] = weights_dir.name --> 212 proto, mil_program = mil_convert_to_proto( 213 model, 214 convert_from, 215 convert_to, 216 registry, 217 **kwargs 218 ) 220 _reset_conversion_state() 222 if convert_to == 'milinternal': File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/converter.py:285, in mil_convert_to_proto(model, convert_from, convert_to, converter_registry, main_pipeline, **kwargs) 280 frontend_pipeline, backend_pipeline = _construct_other_pipelines( 281 main_pipeline, convert_from, convert_to 282 ) 284 frontend_converter = frontend_converter_type() --> 285 prog = frontend_converter(model, **kwargs) 286 PipelineManager.apply_pipeline(prog, frontend_pipeline) 288 PipelineManager.apply_pipeline(prog, main_pipeline) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/converter.py:108, in TorchFrontend.__call__(self, *args, **kwargs) 105 def __call__(self, *args, **kwargs): 106 from .frontend.torch.load import load --> 108 return load(*args, **kwargs) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py:63, in load(model_spec, inputs, specification_version, debug, outputs, cut_at_symbols, **kwargs) 55 inputs = _convert_to_torch_inputtype(inputs) 56 converter = TorchConverter( 57 torchscript, 58 inputs, (...) 61 specification_version, 62 ) ---> 63 return _perform_torch_convert(converter, debug) File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py:102, in _perform_torch_convert(converter, debug) 100 def _perform_torch_convert(converter, debug): 101 try: --> 102 prog = converter.convert() 103 except RuntimeError as e: 104 if debug and "convert function" in str(e): File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/converter.py:284, in TorchConverter.convert(self) 281 self.convert_const() 283 # Add the rest of the operations --> 284 convert_nodes(self.context, self.graph) 286 graph_outputs = [self.context[name] for name in self.graph.outputs] 288 # An output can be None when it's a None constant, which happens 289 # in Fairseq MT. File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py:88, in convert_nodes(context, graph) 83 raise RuntimeError( 84 "PyTorch convert function for op '{}' not implemented.".format(node.kind) 85 ) 87 context.prepare_for_conversion(node) ---> 88 add_op(context, node) 90 # We've generated all the outputs the graph needs, terminate conversion. 91 if _all_outputs_present(context, graph): File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py:4587, in triu(context, node) 4585 @register_torch_op 4586 def triu(context, node): -> 4587 inputs = _get_inputs(context, node, expected=2) 4588 x = inputs[0] 4589 diagonal = inputs[1] File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py:188, in _get_inputs(context, node, expected, min_expected) 182 def _get_inputs(context, node, expected=None, min_expected=None) -> List[Var]: 183 """ 184 Look up a node's inputs in @context and return them as a list. If 185 @expected is not None, also verifies the number of inputs matches the 186 value of @expected. 187 """ --> 188 inputs = [context[name] for name in node.inputs] 189 if expected is not None: 190 expected = [expected] if not isinstance(expected, (list, tuple)) else expected File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py:188, in (.0) 182 def _get_inputs(context, node, expected=None, min_expected=None) -> List[Var]: 183 """ 184 Look up a node's inputs in @context and return them as a list. If 185 @expected is not None, also verifies the number of inputs matches the 186 value of @expected. 187 """ --> 188 inputs = [context[name] for name in node.inputs] 189 if expected is not None: 190 expected = [expected] if not isinstance(expected, (list, tuple)) else expected File ~/miniconda3/envs/prod/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/converter.py:86, in TranscriptionContext.__getitem__(self, torch_name) 84 if torch_name in current_graph: 85 return self._current_graph[idx][torch_name] ---> 86 raise ValueError( 87 "Torch var {} not found in context {}".format(torch_name, self.name) 88 ) ValueError: Torch var mask.3 not found in context ```