Lightning-AI/pytorch-lightning

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

Open

#16.367 geöffnet am 15. Jan. 2023

Auf GitHub ansehen
 (1 Kommentar) (0 Reaktionen) (0 zugewiesene Personen)Python (3.233 Forks)batch import
help wantedquestion

Repository-Metriken

Stars
 (26.687 Stars)
PR-Merge-Metriken
 (Durchschn. Merge 9T 15h) (3 gemergte PRs in 30 T)

Beschreibung

📚 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