|
|
|
|
|
|
|
|
|
import io |
|
import os |
|
|
|
import onnx |
|
import torch |
|
import torch.nn as nn |
|
from onnxsim import simplify as simplify_func |
|
|
|
__all__ = ["export_onnx"] |
|
|
|
|
|
def export_onnx( |
|
model: nn.Module, export_path: str, sample_inputs: any, simplify=True, opset=11 |
|
) -> None: |
|
"""Export a model to a platform-specific onnx format. |
|
|
|
Args: |
|
model: a torch.nn.Module object. |
|
export_path: export location. |
|
sample_inputs: Any. |
|
simplify: a flag to turn on onnx-simplifier |
|
opset: int |
|
""" |
|
model.eval() |
|
|
|
buffer = io.BytesIO() |
|
with torch.no_grad(): |
|
torch.onnx.export(model, sample_inputs, buffer, opset_version=opset) |
|
buffer.seek(0, 0) |
|
if simplify: |
|
onnx_model = onnx.load_model(buffer) |
|
onnx_model, success = simplify_func(onnx_model) |
|
assert success |
|
new_buffer = io.BytesIO() |
|
onnx.save(onnx_model, new_buffer) |
|
buffer = new_buffer |
|
buffer.seek(0, 0) |
|
|
|
if buffer.getbuffer().nbytes > 0: |
|
save_dir = os.path.dirname(export_path) |
|
os.makedirs(save_dir, exist_ok=True) |
|
with open(export_path, "wb") as f: |
|
f.write(buffer.read()) |
|
|
