-
Notifications
You must be signed in to change notification settings - Fork 5
/
run.py
249 lines (197 loc) · 9.9 KB
/
run.py
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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2024, Ampere Computing LLC
try:
from utils import misc # noqa
except ModuleNotFoundError:
import os
import sys
filename = "set_env_variables.sh"
directory = os.path.realpath(__file__).split("/")[:-1]
for idx in range(1, len(directory) - 1):
subdir = "/".join(directory[:-idx])
if filename in os.listdir(subdir):
print(f"\nPlease run \033[91m'source {os.path.join(subdir, filename)}'\033[0m first.")
break
else:
print(f"\n\033[91mFAIL: Couldn't find {filename}, are you running this script as part of Ampere Model Library?"
f"\033[0m")
sys.exit(1)
def parse_args():
import argparse
parser = argparse.ArgumentParser(description="Run ResNet-50 v1.5 model.")
parser.add_argument("-m", "--model_path",
type=str,
help="path to the model")
parser.add_argument("-p", "--precision",
type=str, choices=["fp32", "bf16", "fp16", "int8"], required=True,
help="precision of the model provided")
parser.add_argument("-b", "--batch_size",
type=int, default=1,
help="batch size to feed the model with")
parser.add_argument("-f", "--framework",
type=str,
choices=["tf", "ort", "pytorch"], required=True,
help="specify the framework in which a model should be run")
parser.add_argument("--timeout",
type=float, default=60.0,
help="timeout in seconds")
parser.add_argument("--num_runs",
type=int,
help="number of passes through network to execute")
parser.add_argument("--images_path",
type=str,
help="path to directory with ImageNet validation images")
parser.add_argument("--labels_path",
type=str,
help="path to file with validation labels")
return parser.parse_args()
def run_tf(model_path, batch_size, num_runs, timeout, images_path, labels_path):
from utils.benchmark import run_model
from utils.cv.imagenet import ImageNet
from utils.tf import TFFrozenModelRunner
def run_single_pass(tf_runner, imagenet):
shape = (224, 224)
tf_runner.set_input_tensor("input_tensor:0", imagenet.get_input_array(shape))
output = tf_runner.run(batch_size)
for i in range(batch_size):
imagenet.submit_predictions(
i,
imagenet.extract_top1(output["softmax_tensor:0"][i]),
imagenet.extract_top5(output["softmax_tensor:0"][i])
)
dataset = ImageNet(batch_size, "RGB", images_path, labels_path,
pre_processing="VGG", is1001classes=True)
runner = TFFrozenModelRunner(model_path, ["softmax_tensor:0"])
return run_model(run_single_pass, runner, dataset, batch_size, num_runs, timeout)
def run_tflite(model_path, batch_size, num_runs, timeout, images_path, labels_path):
from utils.benchmark import run_model
from utils.cv.imagenet import ImageNet
from utils.tflite import TFLiteRunner
def run_single_pass(tflite_runner, imagenet):
shape = (224, 224)
tflite_runner.set_input_tensor(tflite_runner.input_details[0]['index'], imagenet.get_input_array(shape))
tflite_runner.run(batch_size)
output_tensor = tflite_runner.get_output_tensor(tflite_runner.output_details[0]['index'])
for i in range(batch_size):
imagenet.submit_predictions(
i,
imagenet.extract_top1(output_tensor[i]),
imagenet.extract_top5(output_tensor[i])
)
dataset = ImageNet(batch_size, "RGB", images_path, labels_path,
pre_processing="VGG", is1001classes=True)
runner = TFLiteRunner(model_path)
return run_model(run_single_pass, runner, dataset, batch_size, num_runs, timeout)
def run_pytorch_fp(model_name, batch_size, num_runs, timeout, images_path, labels_path, disable_jit_freeze=False):
from utils.benchmark import run_model
from utils.cv.imagenet import ImageNet
from utils.pytorch import PyTorchRunner
import torch
import torchvision
def run_single_pass(pytorch_runner, imagenet):
shape = (224, 224)
output = pytorch_runner.run(batch_size, torch.from_numpy(imagenet.get_input_array(shape))).float()
for i in range(batch_size):
imagenet.submit_predictions(
i,
imagenet.extract_top1(output[i]),
imagenet.extract_top5(output[i])
)
dataset = ImageNet(batch_size, "RGB", images_path, labels_path,
pre_processing='PyTorch', is1001classes=False, order='NCHW')
runner = PyTorchRunner(torchvision.models.__dict__[model_name](pretrained=True),
disable_jit_freeze=disable_jit_freeze)
return run_model(run_single_pass, runner, dataset, batch_size, num_runs, timeout)
def run_pytorch_cuda(model_name, batch_size, num_runs, timeout, images_path, labels_path, disable_jit_freeze=False,
**kwargs):
from utils.benchmark import run_model
from utils.cv.imagenet import ImageNet
from utils.pytorch import PyTorchRunner
import torch
import torchvision
def run_single_pass(pytorch_runner, imagenet):
shape = (224, 224)
output = pytorch_runner.run(batch_size, torch.from_numpy(imagenet.get_input_array(shape)).cuda()).cpu()
for i in range(batch_size):
imagenet.submit_predictions(
i,
imagenet.extract_top1(output[i]),
imagenet.extract_top5(output[i])
)
dataset = ImageNet(batch_size, "RGB", images_path, labels_path,
pre_processing='PyTorch', is1001classes=False, order='NCHW')
runner = PyTorchRunner(torchvision.models.__dict__[model_name](pretrained=True).cuda(),
disable_jit_freeze=True)
return run_model(run_single_pass, runner, dataset, batch_size, num_runs, timeout)
def run_ort_fp(model_path, batch_size, num_runs, timeout, images_path, labels_path):
from utils.benchmark import run_model
from utils.cv.imagenet import ImageNet
from utils.ort import OrtRunner
def run_single_pass(ort_runner, imagenet):
shape = (224, 224)
ort_runner.set_input_tensor("input_tensor:0", imagenet.get_input_array(shape).astype("float16"))
output = ort_runner.run(batch_size)
for i in range(batch_size):
imagenet.submit_predictions(
i,
imagenet.extract_top1(output[0][i]),
imagenet.extract_top5(output[0][i])
)
dataset = ImageNet(batch_size, "RGB", images_path, labels_path,
pre_processing="VGG", is1001classes=True)
runner = OrtRunner(model_path)
return run_model(run_single_pass, runner, dataset, batch_size, num_runs, timeout)
def run_tf_fp32(model_path, batch_size, num_runs, timeout, images_path, labels_path, **kwargs):
return run_tf(model_path, batch_size, num_runs, timeout, images_path, labels_path)
def run_tf_fp16(model_path, batch_size, num_runs, timeout, images_path, labels_path, **kwargs):
return run_tf(model_path, batch_size, num_runs, timeout, images_path, labels_path)
def run_tf_bf16(model_path, batch_size, num_runs, timeout, images_path, labels_path, **kwargs):
return run_tf(model_path, batch_size, num_runs, timeout, images_path, labels_path)
def run_pytorch_fp32(model_name, batch_size, num_runs, timeout, images_path, labels_path, **kwargs):
return run_pytorch_fp(model_name, batch_size, num_runs, timeout, images_path, labels_path)
def run_tflite_int8(model_path, batch_size, num_runs, timeout, images_path, labels_path, **kwargs):
return run_tflite(model_path, batch_size, num_runs, timeout, images_path, labels_path)
def run_ort_fp16(model_path, batch_size, num_runs, timeout, images_path, labels_path, **kwargs):
return run_ort_fp(model_path, batch_size, num_runs, timeout, images_path, labels_path)
def main():
from utils.misc import print_goodbye_message_and_die, download_ampere_imagenet
args = parse_args()
download_ampere_imagenet()
if args.framework == "tf":
if args.model_path is None:
print_goodbye_message_and_die(
"a path to model is unspecified!")
if args.precision == "fp32":
run_tf_fp32(**vars(args))
elif args.precision == "fp16":
run_tf_fp16(**vars(args))
elif args.precision == "bf16":
run_tf_bf16(**vars(args))
elif args.precision == "int8":
run_tflite_int8(**vars(args))
else:
print_goodbye_message_and_die(
"this model seems to be unsupported in a specified precision: " + args.precision)
elif args.framework == "pytorch":
import torch
if torch.cuda.is_available():
run_pytorch_cuda(model_name="resnet50", **vars(args))
elif args.precision == "fp32":
run_pytorch_fp32(model_name="resnet50", **vars(args))
else:
print_goodbye_message_and_die(
"this model seems to be unsupported in a specified precision: " + args.precision)
elif args.framework == "ort":
if args.model_path is None:
print_goodbye_message_and_die(
"a path to model is unspecified!")
if args.precision == "fp16":
run_ort_fp16(**vars(args))
else:
print_goodbye_message_and_die(
"this model seems to be unsupported in a specified precision: " + args.precision)
else:
print_goodbye_message_and_die(
"this model seems to be unsupported in a specified framework: " + args.framework)
if __name__ == "__main__":
main()