PyTorch And Torch-MIGraphX

If your model starts in PyTorch, you have two common routes:

PyTorch -> ONNX -> MIGraphX
PyTorch -> Torch-MIGraphX

Choose the route based on how your application deploys models.

ONNX Export Route

Use this when your serving application is C++ or already consumes ONNX:

import torch

model = ...
example = torch.randn(1, 3, 640, 640)

torch.onnx.export(
    model,
    example,
    "model.onnx",
    input_names=["images"],
    output_names=["output"],
    opset_version=17,
    dynamic_axes={"images": {0: "batch"}, "output": {0: "batch"}},
)

Then compile with MIGraphX:

migraphx-driver compile model.onnx --gpu --save model.mxr

Torch-MIGraphX Route

Use Torch-MIGraphX when you want to stay inside a PyTorch-first inference path. The exact API depends on the Torch-MIGraphX version paired with your ROCm release, so keep this route tied to the upstream Torch-MIGraphX documentation for your installed version.

Recommendation For Services

For a C++ service or video pipeline, prefer:

PyTorch training/export -> ONNX artifact -> MIGraphX compile -> saved program

This gives you a stable artifact boundary between training and serving.