From 0456979ded1098b122f0f1be49f0402b3735f4e0 Mon Sep 17 00:00:00 2001 From: Zhao Shenyang Date: Tue, 14 Feb 2023 08:00:22 +0800 Subject: [PATCH 1/3] feature(frameworks): bentoml.onnx runner accept kwargs --- .github/workflows/frameworks.yml | 3 +- src/bentoml/_internal/frameworks/onnx.py | 48 ++++++++--- tests/integration/frameworks/models/onnx.py | 91 ++++++++++++++++++++- 3 files changed, 127 insertions(+), 15 deletions(-) diff --git a/.github/workflows/frameworks.yml b/.github/workflows/frameworks.yml index ea431a17cb9..536d891e934 100644 --- a/.github/workflows/frameworks.yml +++ b/.github/workflows/frameworks.yml @@ -96,6 +96,7 @@ jobs: - *related - src/bentoml/onnx.py - src/bentoml/_internal/frameworks/onnx.py + - src/bentoml/_internal/frameworks/utils/onnx.py - tests/integration/frameworks/models/onnx.py picklable_model: - *related @@ -498,7 +499,7 @@ jobs: - name: Install dependencies run: | pip install . - pip install onnx onnxruntime skl2onnx + pip install onnx onnxruntime skl2onnx transformers[onnx] pip install -r requirements/tests-requirements.txt - name: Run tests and generate coverage report diff --git a/src/bentoml/_internal/frameworks/onnx.py b/src/bentoml/_internal/frameworks/onnx.py index 508d1a7d9df..cc489be5ed9 100644 --- a/src/bentoml/_internal/frameworks/onnx.py +++ b/src/bentoml/_internal/frameworks/onnx.py @@ -26,6 +26,7 @@ from .utils.onnx import ONNXArgType from .utils.onnx import ONNXArgCastedType + from .utils.onnx import ONNXArgCastingFuncType ProvidersType = list[str | tuple[str, dict[str, t.Any]]] @@ -65,7 +66,8 @@ class ONNXOptions(ModelOptions): input_specs: dict[str, list[dict[str, t.Any]]] = attr.field(factory=dict) output_specs: dict[str, list[dict[str, t.Any]]] = attr.field(factory=dict) providers: ProvidersType = attr.field(default=None) - session_options: t.Optional["ort.SessionOptions"] = attr.field(default=None) + session_options: "ort.SessionOptions" | None = attr.field(default=None) + use_kwargs_inputs: bool = attr.field(default=False) def get(tag_like: str | Tag) -> bentoml.Model: @@ -407,7 +409,13 @@ def add_runnable_method( output_specs: list[dict[str, t.Any]], ): - casting_funcs = [gen_input_casting_func(spec) for spec in input_specs] + casting_funcs: list[ONNXArgCastingFuncType] = [] + input_name2casting_func: dict[str, ONNXArgCastingFuncType] = {} + + for spec in input_specs: + casting_f: ONNXArgCastingFuncType = gen_input_casting_func(spec) + casting_funcs.append(casting_f) + input_name2casting_func[spec["name"]] = casting_f if len(output_specs) > 1: @@ -419,17 +427,31 @@ def _process_output(outs): def _process_output(outs): return outs[0] - def _run(self: ONNXRunnable, *args: ONNXArgType) -> t.Any: - casted_args = [ - casting_funcs[idx](args[idx]) for idx in range(len(casting_funcs)) - ] - - input_names: dict[str, ONNXArgCastedType] = { - i.name: val for i, val in zip(self.model.get_inputs(), casted_args) - } - output_names: list[str] = [o.name for o in self.model.get_outputs()] - raw_outs = self.predict_fns[method_name](output_names, input_names) - return _process_output(raw_outs) + if options.use_kwargs_inputs: + + def _run(self: ONNXRunnable, **kwargs: ONNXArgType) -> t.Any: + + input_names: dict[str, ONNXArgCastedType] = { + name: input_name2casting_func[name](val) + for name, val in kwargs.items() + } + output_names: list[str] = [o.name for o in self.model.get_outputs()] + raw_outs = self.predict_fns[method_name](output_names, input_names) + return _process_output(raw_outs) + + else: + + def _run(self: ONNXRunnable, *args: ONNXArgType) -> t.Any: + casted_args = [ + casting_funcs[idx](args[idx]) for idx in range(len(casting_funcs)) + ] + + input_names: dict[str, ONNXArgCastedType] = { + i.name: val for i, val in zip(self.model.get_inputs(), casted_args) + } + output_names: list[str] = [o.name for o in self.model.get_outputs()] + raw_outs = self.predict_fns[method_name](output_names, input_names) + return _process_output(raw_outs) ONNXRunnable.add_method( _run, diff --git a/tests/integration/frameworks/models/onnx.py b/tests/integration/frameworks/models/onnx.py index 7c6828f957d..4dd8fc3aa70 100644 --- a/tests/integration/frameworks/models/onnx.py +++ b/tests/integration/frameworks/models/onnx.py @@ -12,6 +12,8 @@ import torch.nn as nn import onnxruntime as ort from skl2onnx import convert_sklearn +from transformers import AutoTokenizer +from transformers import AutoModelForSequenceClassification from sklearn.datasets import load_iris from sklearn.ensemble import RandomForestClassifier from sklearn.preprocessing import LabelEncoder @@ -279,4 +281,91 @@ def _check( onnx_le_models.append(onnx_le_model) -models: list[FrameworkTestModel] = [onnx_pytorch_model, onnx_rf_model] + onnx_le_models +# tiny bert model +TINY_BERT_MODEL_ID = "prajjwal1/bert-tiny" + + +def make_bert_onnx_model(tmpdir) -> tuple[onnx.ModelProto, t.Any]: + model_id = TINY_BERT_MODEL_ID + bert_model = AutoModelForSequenceClassification.from_pretrained(model_id) + tokenizer = AutoTokenizer(model_id) + sample_text = "This is a sample" + sample_input = tokenizer(sample_text, return_tensors="pt") + model_path = os.path.join(tmpdir, "bert-tiny.onnx") + + torch.onnx.export( + bert_model, + tuple(sample_input.values()), + f=model_path, + input_names=["input_ids", "attention_mask", "token_type_ids"], + output_names=["logits"], + dynamic_axes={ + "input_ids": {0: "batch_size", 1: "sequence"}, + "attention_mask": {0: "batch_size", 1: "sequence"}, + "logits": {0: "batch_size", 1: "sequence"}, + }, + do_constant_folding=True, + opset_version=13, + ) + + onnx_model = onnx.load(model_path) + + expected_input = tokenizer(sample_text, return_tensors="np") + model_output = bert_model(**sample_input) + expected_output = model_output.logits.detach().to("cpu").numpy() + expected_data = (expected_input, expected_output) + return (onnx_model, expected_data) + + +onnx_bert_raw_model, _expected_data = make_bert_onnx_model() +bert_input, bert_expected_output = _expected_data + + +def method_caller_kwargs( + framework_test_model: FrameworkTestModel, + method: str, + args: list[t.Any], + kwargs: dict[str, t.Any], +): + with tempfile.NamedTemporaryFile() as temp: + onnx.save(framework_test_model.model, temp.name) + ort_sess = ort.InferenceSession(temp.name, providers=["CPUExecutionProvider"]) + + def to_numpy(item): + if isinstance(item, np.ndarray): + pass + elif isinstance(item, torch.Tensor): + item = item.detach().to("cpu").numpy() + return item + + input_names = {k: list(v) for k, v in kwargs} + output_names = [o.name for o in ort_sess.get_outputs()] + out = getattr(ort_sess, method)(output_names, input_names)[0] + return out + + +onnx_bert_model = FrameworkTestModel( + name="onnx_bert_model", + model=onnx_bert_raw_model, + model_method_caller=method_caller_kwargs, + model_signatures={"run": {"batchable": True}}, + configurations=[ + Config( + test_inputs={ + "run": [ + Input( + input_args=[], + input_kwargs=bert_input, + expected=close_to(bert_expected_output), + ), + ], + }, + load_kwargs={"use_kwargs_inputs": True}, + check_model=check_model, + ), + ], +) + +models: list[FrameworkTestModel] = ( + [onnx_pytorch_model, onnx_rf_model] + onnx_le_models + [onnx_bert_model] +) From 7672e7f9c6d272fa8f6c3a2ce54c7c1047dc318f Mon Sep 17 00:00:00 2001 From: Zhao Shenyang Date: Tue, 14 Feb 2023 15:47:47 +0800 Subject: [PATCH 2/3] fix testing --- tests/integration/frameworks/models/onnx.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/integration/frameworks/models/onnx.py b/tests/integration/frameworks/models/onnx.py index 4dd8fc3aa70..ca44432914d 100644 --- a/tests/integration/frameworks/models/onnx.py +++ b/tests/integration/frameworks/models/onnx.py @@ -288,7 +288,7 @@ def _check( def make_bert_onnx_model(tmpdir) -> tuple[onnx.ModelProto, t.Any]: model_id = TINY_BERT_MODEL_ID bert_model = AutoModelForSequenceClassification.from_pretrained(model_id) - tokenizer = AutoTokenizer(model_id) + tokenizer = AutoTokenizer.from_pretrained(model_id) sample_text = "This is a sample" sample_input = tokenizer(sample_text, return_tensors="pt") model_path = os.path.join(tmpdir, "bert-tiny.onnx") @@ -317,8 +317,9 @@ def make_bert_onnx_model(tmpdir) -> tuple[onnx.ModelProto, t.Any]: return (onnx_model, expected_data) -onnx_bert_raw_model, _expected_data = make_bert_onnx_model() -bert_input, bert_expected_output = _expected_data +with tempfile.TemporaryDirectory() as tmpdir: + onnx_bert_raw_model, _expected_data = make_bert_onnx_model(tmpdir) + bert_input, bert_expected_output = _expected_data def method_caller_kwargs( @@ -341,6 +342,7 @@ def to_numpy(item): input_names = {k: list(v) for k, v in kwargs} output_names = [o.name for o in ort_sess.get_outputs()] out = getattr(ort_sess, method)(output_names, input_names)[0] + print("hahahah lkasjdfklasfdsaf") return out From 607ec46d50474ca6777df1e405da7d8ee51711bd Mon Sep 17 00:00:00 2001 From: Zhao Shenyang Date: Wed, 22 Feb 2023 04:46:51 +0800 Subject: [PATCH 3/3] add pytorch to onnx ci test dependency --- .github/workflows/frameworks.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/frameworks.yml b/.github/workflows/frameworks.yml index 536d891e934..2206f12fe36 100644 --- a/.github/workflows/frameworks.yml +++ b/.github/workflows/frameworks.yml @@ -499,7 +499,7 @@ jobs: - name: Install dependencies run: | pip install . - pip install onnx onnxruntime skl2onnx transformers[onnx] + pip install onnx onnxruntime skl2onnx transformers[onnx] torch pip install -r requirements/tests-requirements.txt - name: Run tests and generate coverage report