Lightning-AI/pytorch-lightning
在 GitHub 查看How to generalize the first dimension (often assigned to batch size) on ONNX production inference?
Open
#16,367 创建于 2023年1月15日
help wantedquestion
仓库指标
- Star
- (26,687 star)
- PR 合并指标
- (平均合并 9天 15小时) (30 天内合并 3 个 PR)
描述
📚 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?