1
1
import argparse
2
2
import importlib
3
3
import re
4
+ import time
4
5
import unittest
5
6
7
+ from datetime import timedelta
6
8
from typing import Any
7
9
8
10
import torch
@@ -43,6 +45,7 @@ def run_test( # noqa: C901
43
45
"""
44
46
45
47
error_statistics : list [ErrorStatistics ] = []
48
+ extra_stats = {}
46
49
47
50
# Helper method to construct the summary.
48
51
def build_result (
@@ -57,6 +60,7 @@ def build_result(
57
60
result = result ,
58
61
error = error ,
59
62
tensor_error_statistics = error_statistics ,
63
+ ** extra_stats ,
60
64
)
61
65
62
66
# Ensure the model can run in eager mode.
@@ -71,11 +75,16 @@ def build_result(
71
75
return build_result (TestResult .UNKNOWN_FAIL , e )
72
76
73
77
if flow .quantize :
78
+ start_time = time .perf_counter ()
74
79
try :
75
80
tester .quantize (
76
81
flow .quantize_stage_factory () if flow .quantize_stage_factory else None
77
82
)
83
+ elapsed = time .perf_counter () - start_time
84
+ extra_stats ["quantize_time" ] = timedelta (seconds = elapsed )
78
85
except Exception as e :
86
+ elapsed = time .perf_counter () - start_time
87
+ extra_stats ["quantize_time" ] = timedelta (seconds = elapsed )
79
88
return build_result (TestResult .QUANTIZE_FAIL , e )
80
89
81
90
try :
@@ -86,9 +95,14 @@ def build_result(
86
95
except Exception as e :
87
96
return build_result (TestResult .EXPORT_FAIL , e )
88
97
98
+ lower_start_time = time .perf_counter ()
89
99
try :
90
100
tester .to_edge_transform_and_lower ()
101
+ elapsed = time .perf_counter () - lower_start_time
102
+ extra_stats ["lower_time" ] = timedelta (seconds = elapsed )
91
103
except Exception as e :
104
+ elapsed = time .perf_counter () - lower_start_time
105
+ extra_stats ["lower_time" ] = timedelta (seconds = elapsed )
92
106
return build_result (TestResult .LOWER_FAIL , e )
93
107
94
108
is_delegated = any (
@@ -183,6 +197,7 @@ def parse_args():
183
197
"--report" ,
184
198
nargs = "?" ,
185
199
help = "A file to write the test report to, in CSV format." ,
200
+ default = "backend_test_report.csv" ,
186
201
)
187
202
return parser .parse_args ()
188
203
0 commit comments