From c18845a2fdfbf30b068c452962217393b9703886 Mon Sep 17 00:00:00 2001 From: YdrMaster Date: Mon, 27 Feb 2023 10:03:06 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=E6=8E=A8=E7=90=86?= =?UTF-8?q?=E6=8E=A5=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: YdrMaster --- pyinfinitensor/src/pyinfinitensor/onnx.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/pyinfinitensor/src/pyinfinitensor/onnx.py b/pyinfinitensor/src/pyinfinitensor/onnx.py index ffb61529..1f6afe19 100644 --- a/pyinfinitensor/src/pyinfinitensor/onnx.py +++ b/pyinfinitensor/src/pyinfinitensor/onnx.py @@ -572,6 +572,24 @@ class OnnxStub: return ctx.build(name) + def init(self) -> None: + self.handler.data_malloc() + + def run(self) -> None: + self.handler.run() + + def put_int32(self, name: str) -> None: + self.handler.copy_int32(self.inputs[name]) + + def put_int64(self, name: str) -> None: + self.handler.copy_int64(self.inputs[name]) + + def put_float(self, name: str) -> None: + self.handler.copy_float(self.inputs[name]) + + def take_float(self) -> List[float]: + return next(self.handler.outputs.values()).copyFloats() + def from_onnx(model: ModelProto, runtime): stub = OnnxStub(model, runtime)