使用 torch-mlir 加载 MLIR 文件作为模型

Question

我已经使用 torch-mlir 将火炬模型转换为 MLIR，如 github 示例中所示：

resnet18 = torchvision.models.resnet18(pretrained=True)
resnet18.eval()
module = torch_mlir.compile(resnet18, torch.ones(1, 3, 224, 224), output_type="torch")
print("TORCH OutputType\n", 
module.operation.get_asm(large_elements_limit=10))

在那之后，我改变了 MLIR，现在，我想做相反的事情，将这个 MLIR 作为一个模块加载到我的 python 代码中，继续像示例中那样编译它：

backend = refbackend.RefBackendLinalgOnTensorsBackend()
compiled = backend.compile(module)
jit_module = backend.load(compiled)
predictions(resnet18.forward, jit_module.forward, img, labels)

我只需要两个代码之间的一行来加载 MLIR，但我在互联网上找不到任何相关信息。有人知道怎么做吗？

Answer 1

import torch_mlir.ir

src = open("file.mlir", "r").read()

with torch_mlir.ir.Context() as ctx:
  torch_mlir.dialects.torch.register_dialect(ctx)
  with torch_mlir.ir.Location.unknown() as loc:
    module = torch_mlir.ir.Module.parse(src)
  module.operation.print()