PyTorch转NNEF模型转换器

Quickstart

pip install torch_to_nnef   # needs pip >= 23.2

import torch
from torch_to_nnef import export_model_to_nnef, TractNNEF

model = torch.nn.Linear(8, 4).eval()
example_input = torch.randn(1, 8)

export_model_to_nnef(
    model=model,
    args=example_input,
    file_path_export="my_model.nnef.tgz",
    # version pins the tract opset to target; check_io downloads that tract
    # build once and asserts its output matches PyTorch on example_input.
    inference_target=TractNNEF(version="0.21.13", check_io=True),
    input_names=["input"],
    output_names=["output"],
)

This writes a self-contained my_model.nnef.tgz; with check_io=True the call also raises if tract and PyTorch outputs diverge, so reaching the end means the export is validated. The getting started tutorial (linked below) walks through running the archive in tract from Python, Rust, and the CLI.

<details> <summary>What's inside <code>my_model.nnef.tgz</code></summary>

The archive keeps the graph and the weights apart:

graph.nnef weight.dat bias.dat

and graph.nnef is plain text you can read end to end (metadata header elided):

graph network(input) -> (output) { input = tract_core_external(shape = [1, 8], datum_type = 'f32'); weight = variable<scalar>(label = 'weight', shape = [4, 8]); bias = variable<scalar>(label = 'bias', shape = [4]); bias_aligned_rank_expanded = unsqueeze(bias, axes = [0]); output = linear(input, weight, bias_aligned_rank_expanded); }

</details>