Lightning-AI/pytorch-lightning

How to generalize the first dimension (often assigned to batch size) on ONNX production inference?

Open

#16,367 opened on Jan 15, 2023

View on GitHub
 (1 comment) (0 reactions) (0 assignees)Python (3,233 forks)batch import
help wantedquestion

Repository metrics

Stars
 (26,687 stars)
PR merge metrics
 (Avg merge 9d 15h) (3 merged PRs in 30d)

Description

📚 Documentation

Since we have to pass input_sample (torch.randn((1, 64), for instance), when compiling the PL model to ONNX, we must pass only input with shape=(1,64), on inference time:

# compiling the PL model to ONNX
model = PLModel()
filepath = "model.onnx"
input_sample = torch.randn((1, 64))
model.to_onnx(filepath, input_sample, export_params=True)

# inference time
import onnxruntime

ort_session = onnxruntime.InferenceSession(filepath)
input_name = ort_session.get_inputs()[0].name
ort_inputs = {input_name: np.random.randn(1, 64)}
ort_outs = ort_session.run(None, ort_inputs)

Then, how to generalize the first dimension (often assigned to batch size) on onnx production inference?

Contributor guide

Research direction
Check the PyTorch Lightning documentation on how to export a model to ONNX with dynamic axes. The key is to use the 'dynamic axes' parameter in torch.onnx.export. Then verify that the ONNX model accepts variable batch sizes during inference with ONNX Runtime.
Tech stack
pythonpytorch
Domain
backendapi
Issue type
Documentation
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
2
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-3 hours
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
Fresh
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
Clear
Prerequisites
PythonPyTorch
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
75