Lightning-AI / Lightning-AI/pytorch-lightning

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

Open
#16,367 1 comment 0 reactions 0 assignees View on GitHub
help wanted question
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

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:

```python
# 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

Open the contributing guide

Research direction

Start with the issue's model.to_onnx example and the ONNX Runtime InferenceSession example to identify the documentation entry point for production inference. Document how the first input dimension should be handled for variable batch sizes, including the expected export and inference behavior, and verify the example remains consistent.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
documentation, machine-learning
Issue type
Documentation
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.