llvm / torch-mlir

The Torch-MLIR project aims to provide first class support from the PyTorch ecosystem to the MLIR ecosystem.
Other
1.32k stars 487 forks source link

Support lowering 'densenet161' from 'torchvision' to 'torch` dialect #2950

Open terfendail opened 7 months ago

terfendail commented 7 months ago

Lowering of DenseNet model fails with error.

import torch
import torch_mlir
import torchvision
mdl = torchvision.models.densenet.densenet161()
shp = torch_mlir.TensorPlaceholder([1, 3, 224, 224], torch.float32)
mlir = torch_mlir.compile(mdl, shp, "torch")

Execution of the script cause following error:

File "/pytorch_venv/lib/python3.11/site-packages/torch_mlir/__init__.py", line 458, in compile
run_pipeline_with_repro_report(
File "/pytorch_venv/lib/python3.11/site-packages/torch_mlir/compiler_utils.py", line 73, in run_pipeline_with_repro_report
raise TorchMlirCompilerError(trimmed_message) from None
torch_mlir.compiler_utils.TorchMlirCompilerError: Lowering TorchScript IR -> Torch Backend IR failed with the following diagnostics:
python exception: Failure while executing pass pipeline:
error: unknown: unsupported by backend contract: module initializers
terfendail commented 7 months ago

If the model is loaded inside of the wrapper

import torch
import torch_mlir
import torchvision
class ModelWrapper(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.model = torchvision.models.densenet.densenet161()
        self.model.eval()
    def forward(self, data):
        return self.model(data)[0]

mdl = ModelWrapper()
mdl.eval()
shp = torch_mlir.TensorPlaceholder([-1, 3, 224, 224], torch.float32)
mlir = torch_mlir.compile(mdl, shp, "torch")

the error changes to

Traceback (most recent call last):
  File "/pytorch_venv/convert_test.py", line 24, in <module>
    mlir = torch_mlir.compile(mdl, shp, "torch")
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pytorch_venv/lib/python3.11/site-packages/torch_mlir/__init__.py", line 458, in compile
    run_pipeline_with_repro_report(
  File "/pytorch_venv/lib/python3.11/site-packages/torch_mlir/compiler_utils.py", line 73, in run_pipeline_with_repro_report
    raise TorchMlirCompilerError(trimmed_message) from None
torch_mlir.compiler_utils.TorchMlirCompilerError: Lowering TorchScript IR -> Torch Backend IR failed with the following diagnostics:

python exception: Failure while executing pass pipeline:
error: "aten::select"("/pytorch_venv/convert_test.py":19:15): unsupported by backend contract: tensor with unknown rank
note: "aten::select"("pytorch_venv/convert_test.py":19:15): see current operation: %844 = "torch.aten.squeeze.dim"(%843, %0) : (!torch.vtensor<[?,1000],f32>, !torch.int) -> !torch.vtensor<*,f32>
note: "aten::select"("/pytorch_venv/convert_test.py":19:15): this is likely due to a missing transfer function in abstract_interp_lib_gen.py