Lightning-AI/pytorch-lightning

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

Open

#16 367 ouverte le 15 janv. 2023

Voir sur GitHub
 (1 commentaire) (0 réactions) (0 assignés)Python (3 233 forks)batch import
help wantedquestion

Métriques du dépôt

Stars
 (26 687 stars)
Métriques de merge PR
 (Merge moyen 9j 15h) (3 PRs mergées en 30 j)

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?

Guide contributeur