ONNX PyTorch

flytekitplugins-onnxpytorch

FlytekitModel Servingonnxpytorchmodel

This plugin allows you to generate ONNX models from your PyTorch models.

Install

pip install flytekitplugins-onnxpytorch

Quick Start(example, may need adjustment)

pip install flytekitplugins-onnxpytorch

from flytekit import task, workflow
from flytekitplugins.onnxpytorch import PyTorch2ONNX, PyTorch2ONNXConfig

config = PyTorch2ONNX(...)

@task
def my_task() -> None:
    ...

@workflow
def my_workflow() -> None:
    my_task()

Available Imports (2)

configPyTorch2ONNX

PyTorch2ONNXConfig is the config used during the pytorch to ONNX conversion.

extends dataclass — configuration or data structure for plugin setup

from flytekitplugins.onnxpytorch import PyTorch2ONNX

configPyTorch2ONNXConfig

PyTorch2ONNXConfig is the config used during the pytorch to ONNX conversion.

extends dataclass — configuration or data structure for plugin setup

from flytekitplugins.onnxpytorch import PyTorch2ONNXConfig

Dependencies

torchonnx-pytorchnetworkxonnxscriptnumpy

Related Plugins

ONNX TensorFlow

Model Serving

This plugin allows you to generate ONNX models from your TensorFlow Keras models.

ONNX ScikitLearn

Model Serving

This plugin allows you to generate ONNX models from your ScikitLearn models.

Dgxc-lepton

Model Serving

A professional Flytekit plugin that enables seamless deployment and management of AI inference endpoints using Lepton AI infrastructure within Flyte workflows.