Lightning-AI/pytorch-lightning

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

Open

#16.367 aperta il 15 gen 2023

Vedi su GitHub
 (1 commento) (0 reazioni) (0 assegnatari)Python (3233 fork)batch import
help wantedquestion

Metriche repository

Star
 (26.687 star)
Metriche merge PR
 (Merge medio 9g 15h) (3 PR mergiate in 30 g)

Descrizione

📚 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?

Guida contributor