Failing to convert the model to MLIR using torch_mlir.compile

[RiRi114]

Hi, I'm trying to use torch-mlir to compile two models, one is a pytorch yolov5 model from https://github.com/ultralytics/yolov5 and the other is wide_and_deep model from https://github.com/jrzaurin/pytorch-widedeep. But I'm having some issues now, when compiling these two models.

YOLOV5

1. python script

import torch import torch_mlir model = torch.hub.load("ultralytics/yolov5", "yolov5s") module = torch_mlir.compile(model, torch.zeros(1,3,640,640), output_type="linalg-on-tensors")

2. error

Traceback (most recent call last):
  File "/home/project/iree-demo/pytorch/test-case-generate/yolo/yolo.py", line 32, in <module>
    module = torch_mlir.compile(yolomodel, torch.zeros(1,3,640,640), output_type="linalg-on-tensors")
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch_mlir/__init__.py", line 404, in compile
    scripted = torch.jit.script(model)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_script.py", line 1284, in script
    return torch.jit._recursive.create_script_module(
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 480, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct
    init_fn(script_module)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct
    init_fn(script_module)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct
    init_fn(script_module)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct
    init_fn(script_module)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 546, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 397, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_script.py", line 1466, in _recursive_compile_class
    return _compile_and_register_class(obj, rcb, _qual_name)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/_recursive.py", line 49, in _compile_and_register_class
    ast = get_jit_class_def(obj, obj.__name__)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/frontend.py", line 234, in get_jit_class_def
    method_defs = [
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/frontend.py", line 235, in <listcomp>
    get_jit_def(obj, name, self_name=self_name, is_classmethod=is_classmethod(obj))
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/frontend.py", line 297, in get_jit_def
    return build_def(parsed_def.ctx, fn_def, type_line, def_name, self_name=self_name, pdt_arg_types=pdt_arg_types)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/frontend.py", line 335, in build_def
    param_list = build_param_list(ctx, py_def.args, self_name, pdt_arg_types)
  File "/root/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/jit/frontend.py", line 363, in build_param_list
    raise NotSupportedError(ctx_range, _vararg_kwarg_err)
torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/root/anaconda3/envs/yolo/lib/python3.10/warnings.py", line 477
    def __exit__(self, *exc_info):
                       ~~~~~~~~~ <--- HERE
        if not self._entered:
            raise RuntimeError("Cannot exit %r without entering first" % self)
'__torch__.warnings.catch_warnings' is being compiled since it was called from 'SPPF.forward'
  File "/root/.cache/torch/hub/ultralytics_yolov5_master/models/common.py", line 229
    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
             ~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
            y1 = self.m(x)

Wide And Deep

1. python script ` import pandas as pd import numpy as np import torch from sklearn.model_selection import train_test_split

from pytorch_widedeep import Trainer from pytorch_widedeep.preprocessing import WidePreprocessor, TabPreprocessor from pytorch_widedeep.models import Wide, TabMlp, WideDeep from pytorch_widedeep.metrics import Accuracy from pytorch_widedeep.datasets import load_adult

import torch_mlir

df = load_adult(as_frame=True) df["income_label"] = (df["income"].apply(lambda x: ">50K" in x)).astype(int) df.drop("income", axis=1, inplace=True) df_train, df_test = train_test_split(df, test_size=0.2, stratify=df.income_label)

Define the 'column set up'

wide_cols = [ "education", "relationship", "workclass", "occupation", "native-country", "gender", ] crossed_cols = [("education", "occupation"), ("native-country", "occupation")]

cat_embed_cols = [ "workclass", "education", "marital-status", "occupation", "relationship", "race", "gender", "capital-gain", "capital-loss", "native-country", ] continuous_cols = ["age", "hours-per-week"] target = "income_label" target = df_train[target].values

prepare the data

wide_preprocessor = WidePreprocessor(wide_cols=wide_cols, crossed_cols=crossed_cols) X_wide = wide_preprocessor.fit_transform(df_train)

tab_preprocessor = TabPreprocessor( cat_embed_cols=cat_embed_cols, continuous_cols=continuous_cols # type: ignore[arg-type] ) X_tab = tab_preprocessor.fit_transform(df_train)

build the model

wide = Wide(input_dim=np.unique(X_wide).shape[0], pred_dim=1) tab_mlp = TabMlp( column_idx=tab_preprocessor.column_idx, cat_embed_input=tab_preprocessor.cat_embed_input, continuous_cols=continuous_cols, ) model = WideDeep(wide=wide, deeptabular=tab_mlp) print("dim: ",np.unique(X_wide).shape[0])

train and validate

trainer = Trainer(model, objective="binary", metrics=[Accuracy]) trainer.fit( X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=5, batch_size=256, )

predict on test

X_wide_te = wide_preprocessor.transform(df_test) X_tab_te = tab_preprocessor.transform(df_test) preds = trainer.predict(X_wide=X_wide_te, X_tab=X_tab_te)

Save and load

Option 1: this will also save training history and lr history if the

LRHistory callback is used

trainer.save(path="model_weights", save_state_dict=True)

Option 2: save as any other torch model

torch.save(model.state_dict(), "model_weights/wd_model.pt")

From here in advance, Option 1 or 2 are the same. I assume the user has

prepared the data and defined the new model components:

1. Build the model

model_new = WideDeep(wide=wide, deeptabular=tab_mlp) model_new.load_state_dict(torch.load("model_weights/wd_model.pt"))

2. Instantiate the trainer

trainer_new = Trainer(model_new, objective="binary")

3. Either start the fit or directly predict

preds = trainer_new.predict(X_wide=X_wide, X_tab=X_tab)

data = [] data.append(torch.from_numpy(X_wide)) data.append(torch.from_numpy(X_tab)) print(model_new)

module = torch_mlir.compile(model_new, data, output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, use_tracing=False) ` 2. Error

Traceback (most recent call last): File "/home/project/iree-demo/pytorch/test-case-generate/widedeep/wide_deep.py", line 105, in module = torch_mlir.compile(model_new, data, output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, use_tracing=False) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch_mlir/init.py", line 404, in compile scripted = torch.jit.script(model) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 1284, in script return torch.jit._recursive.create_script_module( File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 480, in create_script_module return create_script_module_impl(nn_module, concrete_type, stubs_fn) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct init_fn(script_module) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct init_fn(script_module) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct init_fn(script_module) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct init_fn(script_module) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 546, in create_script_module_impl create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs) File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 397, in create_methods_and_properties_from_stubs concrete_type.create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults) RuntimeError: Unable to extract string literal index. ModuleDict indexing is only supported with string literals. For example, 'i = "a"; self.layersi' will fail because i is not a literal. Enumeration of ModuleDict is supported, e.g. 'for k, v in self.items(): out = v(inp)': File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/pytorch_widedeep/models/tabular/embeddings_layers.py", line 173 def forward(self, X: Tensor) -> Tensor: embed = [ self.embed_layers["emb_layer" + self.embed_layers_names[col]]( ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE X[:, self.column_idx[col]].long() )

[ramiro050]

Both errors seem to be from PyTorch, not torch-mlir. When you call torch_mlir.compile, torch-mlir calls torch.jit.script on the module before importing it into torch-mlir. The first error is due to torch.jit.script not supporting functions with variable number of arguments. The second error also seems related to torch.jit.script

[RiRi114]

I understand, thanks for your reply

[DiviyamPathak]

@RiRi114 hello did you find any solution for this. I am trying to do the same