2023-02-02 21:29:47 +08:00
|
|
|
import torch
|
|
|
|
import torchvision.models as models
|
2024-02-07 11:07:59 +08:00
|
|
|
from torch_mlir import torchscript
|
2023-02-02 21:29:47 +08:00
|
|
|
|
|
|
|
model = models.resnet18(pretrained=True)
|
|
|
|
model.eval()
|
2024-04-28 05:16:31 +08:00
|
|
|
data = torch.randn(2, 3, 200, 200)
|
2023-02-02 21:29:47 +08:00
|
|
|
out_stablehlo_mlir_path = "./resnet18_stablehlo.mlir"
|
|
|
|
|
2024-04-28 05:16:31 +08:00
|
|
|
module = torchscript.compile(
|
|
|
|
model, data, output_type=torchscript.OutputType.STABLEHLO, use_tracing=False
|
|
|
|
)
|
2023-02-02 21:29:47 +08:00
|
|
|
with open(out_stablehlo_mlir_path, "w", encoding="utf-8") as outf:
|
2024-05-11 00:28:58 +08:00
|
|
|
outf.write(module.operation.get_asm())
|
2023-02-02 21:29:47 +08:00
|
|
|
|
|
|
|
print(f"StableHLO IR of resent18 successfully written into {out_stablehlo_mlir_path}")
|