1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
|
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import os
import tempfile
import unittest
import numpy as np
import onnx
from onnxscript import proto2python, script
from onnxscript.onnx_opset import opset17 as op
from onnxscript.onnx_types import FLOAT
class TestConverter(unittest.TestCase):
def test_external_tensor(self):
weight = np.random.rand(1024, 10).astype(np.float32)
bias = np.random.rand(10).astype(np.float32)
@script()
def TestFun(X: FLOAT[1024]) -> FLOAT[1024]:
return op.MatMul(X, weight) + bias
model = TestFun.to_model_proto()
with tempfile.TemporaryDirectory() as dir:
# Convert model to use external-tensors and save
modelfile = os.path.join(dir, "model.onnx")
onnx.save_model(
model,
modelfile,
save_as_external_data=True,
all_tensors_to_one_file=False,
size_threshold=32,
convert_attribute=True,
)
# Convert model to python:
pymodel = proto2python(model)
self.assertIn(
"external_tensor('weight', 1, [1024, 10], 'weight', length=40960)", pymodel
)
self.assertIn("external_tensor('bias', 1, [10], 'bias', length=40)", pymodel)
if __name__ == "__main__":
unittest.main()
|
