Lightning-AI / Lightning-AI/pytorch-lightning
How to generalize the first dimension (often assigned to batch size) on ONNX production inference?
- 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
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