onnx-array-api/_unittests/ut_ort/test_ort_profile.py at 4c42a4be846faa3c41a15cab8e44e9c446c0d4a2 · sdpython/onnx-array-api · GitHub

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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import unittest
import os
import numpy as np
from pandas import DataFrame, read_excel
from onnx_array_api.npx import absolute, jit_onnx
from onnx_array_api.ext_test_case import ExtTestCase
from onnx_array_api.ort.ort_optimizers import ort_optimized_model
from onnx_array_api.ort.ort_profile import ort_profile, merge_ort_profile
from onnxruntime.capi._pybind_state import (
    OrtValue as C_OrtValue,
    OrtDevice as C_OrtDevice,
)


class TestOrtProfile(ExtTestCase):
    def test_ort_profile(self):
        def l1_loss(x, y):
            return absolute(x - y).sum()

        def l2_loss(x, y):
            return ((x - y) ** 2).sum()

        def myloss(x, y):
            return l1_loss(x[:, 0], y[:, 0]) + l2_loss(x[:, 1], y[:, 1])

        jitted_myloss = jit_onnx(myloss)
        x = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
        y = np.array([[0.11, 0.22], [0.33, 0.44]], dtype=np.float32)
        jitted_myloss(x, y)
        onx = jitted_myloss.get_onnx()
        feeds = {"x0": x, "x1": y}
        self.assertRaise(lambda: ort_optimized_model(onx, "NO"), ValueError)
        optimized = ort_optimized_model(onx)
        prof = ort_profile(optimized, feeds)
        self.assertIsInstance(prof, DataFrame)
        prof = ort_profile(optimized, feeds, as_df=False)
        self.assertIsInstance(prof, list)

    def test_ort_profile_first_it_out(self):
        def l1_loss(x, y):
            return absolute(x - y).sum()

        def l2_loss(x, y):
            return ((x - y) ** 2).sum()

        def myloss(x, y):
            return l1_loss(x[:, 0], y[:, 0]) + l2_loss(x[:, 1], y[:, 1])

        jitted_myloss = jit_onnx(myloss)
        x = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
        y = np.array([[0.11, 0.22], [0.33, 0.44]], dtype=np.float32)
        jitted_myloss(x, y)
        onx = jitted_myloss.get_onnx()
        feeds = {"x0": x, "x1": y}
        self.assertRaise(lambda: ort_optimized_model(onx, "NO"), ValueError)
        optimized = ort_optimized_model(onx)
        prof = ort_profile(optimized, feeds)
        events = {
            "kernel_time",
            "SequentialExecutor::Execute",
            "model_run",
            "model_loading_array",
            "session_initialization",
        }
        self.assertEqual(set(prof["event_name"]), events)
        agg = ort_profile(optimized, feeds, first_it_out=True, agg=True)
        self.assertIsInstance(agg, DataFrame)
        self.assertLess(agg.shape[0], prof.shape[0])
        self.assertEqual(set(agg.reset_index(drop=False)["event_name"]), events)
        agg = ort_profile(
            optimized, feeds, first_it_out=True, agg=True, agg_op_name=False
        )
        self.assertIsInstance(agg, DataFrame)
        self.assertLess(agg.shape[0], prof.shape[0])
        self.assertEqual(set(agg.reset_index(drop=False)["event_name"]), events)

    def test_ort_profile_ort_value(self):
        def to_ort_value(m):
            device = C_OrtDevice(C_OrtDevice.cpu(), C_OrtDevice.default_memory(), 0)
            ort_value = C_OrtValue.ortvalue_from_numpy(m, device)
            return ort_value

        def l1_loss(x, y):
            return absolute(x - y).sum()

        def l2_loss(x, y):
            return ((x - y) ** 2).sum()

        def myloss(x, y):
            return l1_loss(x[:, 0], y[:, 0]) + l2_loss(x[:, 1], y[:, 1])

        jitted_myloss = jit_onnx(myloss)
        x = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
        y = np.array([[0.11, 0.22], [0.33, 0.44]], dtype=np.float32)
        jitted_myloss(x, y)
        onx = jitted_myloss.get_onnx()
        np_feeds = {"x0": x, "x1": y}
        feeds = {k: to_ort_value(v) for k, v in np_feeds.items()}

        self.assertRaise(lambda: ort_optimized_model(onx, "NO"), ValueError)
        optimized = ort_optimized_model(onx)
        prof = ort_profile(optimized, feeds)
        self.assertIsInstance(prof, DataFrame)
        prof = ort_profile(optimized, feeds, as_df=False)
        self.assertIsInstance(prof, list)

    def test_merge_ort_profile(self):
        data = os.path.join(os.path.dirname(__file__), "data")
        df1 = read_excel(os.path.join(data, "prof_base.xlsx"))
        df2 = read_excel(os.path.join(data, "prof_opti.xlsx"))
        merged, gr = merge_ort_profile(df1, df2)
        self.assertEqual(merged.shape, (23, 9))
        self.assertEqual(
            list(merged.columns),
            [
                "args_op_name",
                "args_output_type_shape",
                "args_input_type_shape",
                "args_provider",
                "idx",
                "durbase",
                "countbase",
                "duropti",
                "countopti",
            ],
        )
        self.assertEqual(gr.shape, (19, 4))
        self.assertEqual(
            list(gr.columns), ["durbase", "duropti", "countbase", "countopti"]
        )


if __name__ == "__main__":
    unittest.main(verbosity=2)