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