加载一个MLIR文件作为模型,使用torch-mlir。

huangapple go评论58阅读模式
英文:

load an MLIR file as a model using torch-mlir

问题

I just need a line in the second code to load the MLIR as a module but I couldn't find anything about it on the internet. Does anyone know how to do it?

英文:

I have translated a torch model into MLIR using torch-mlir as in the github example:

def save_module():
    resnet18 = models.resnet18(pretrained=True)
    resnet18.eval()    
    module = torch_mlir.compile(resnet18, torch.ones(1, 3, 224, 224), output_type="torch")
    open("resnet18torch.mlir", "w").write(str(module))

After that, I changed the MLIR and now, I would like to do the opposite and load this MLIR as a module in my python code to continue to compile it like that:

src = open("resnet18torch.mlir", "r").read()
#transform src to module
backend = refbackend.RefBackendLinalgOnTensorsBackend()
compiled = backend.compile(module)
jit_module = backend.load(compiled)
predictions(resnet18.forward, jit_module.forward, img, labels)

I just need a line in the second code to load the MLIR as a module but I couldn't find anything about it on the internet. Does anyone know how to do it?

答案1

得分: 1

我终于弄清楚了如何加载 torch mlir 并如何在 LLVM Discord 上的某人的帮助下重新使用它。以下是使用 torch_mlir 示例(torchscript_resnet18.py)的用法:

import torch_mlir

def load_module():
    # 加载 torch mlir
    src = open("resnet18torch.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)
    # 将 torch-mlir 转换为 linalg-on-tensors 方言
    torch_mlir.run_pipeline_with_repro_report(module, "builtin.module(torch-backend-to-linalg-on-tensors-backend-pipeline)", "将 Torch 后端 IR 转换为 Linalg-on-Tensors 后端 IR")

    # 编译模块
    backend = refbackend.RefBackendLinalgOnTensorsBackend()
    compiled = backend.compile(module)
    jit_module = backend.load(compiled)
    return jit_module

# 保存模块()
jit_module = load_module()
predictions(jit_module.forward, img, labels)

请注意,我已经将代码中的 HTML 实体(")替换为相应的引号。

英文:

I finally figured out how to load a torch mlir and how to reuse it with the help of someone on the LLVM discord. Here is the usage with the torch_mlir example (torchscript_resnet18.py):

import torch_mlir

def load_module():
    #load the torch mlir
    src = open("resnet18torch.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)
    #translate the torch-mlir to the linalg-on-tensors dialect
    torch_mlir.run_pipeline_with_repro_report(module, "builtin.module(torch-backend-to-linalg-on-tensors-backend-pipeline)", "Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR")

    #print(module)
    #compile the module
    backend = refbackend.RefBackendLinalgOnTensorsBackend()
    compiled = backend.compile(module)
    jit_module = backend.load(compiled)
    return jit_module

#save_module()
jit_module = load_module()
predictions(jit_module.forward, img, labels)

huangapple
  • 本文由 发表于 2023年5月10日 23:18:36
  • 转载请务必保留本文链接:https://go.coder-hub.com/76220147.html
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定