-
Notifications
You must be signed in to change notification settings - Fork 34
/
compute_speeds.py
70 lines (49 loc) · 1.69 KB
/
compute_speeds.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
# ------------------------------------------------------------------------
# Copyright (c) 2021 megvii-model. All Rights Reserved.
# ------------------------------------------------------------------------
# taken from https://gist.github.com/fmassa/c0fbb9fe7bf53b533b5cc241f5c8234c with a few modifications
# taken from detectron2 / fvcore with a few modifications
# https://github.com/facebookresearch/detectron2/blob/master/detectron2/utils/analysis.py
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import numpy as np
import tqdm
import torch
import time
from main import get_args_parser as get_main_args_parser
from models import build_model
from datasets import build_dataset
def warmup(model, inputs, N=10):
for i in range(N):
out = model(inputs)
torch.cuda.synchronize()
def measure_time(model, inputs, N=10):
warmup(model, inputs)
s = time.time()
for i in range(N):
out = model(inputs)
torch.cuda.synchronize()
t = (time.time() - s) / N
return t
def fmt_res(data):
return data.mean(), data.std(), data.min(), data.max()
def benchmark():
main_args = get_main_args_parser().parse_args()
dataset = build_dataset('val', main_args)
model, _, _ = build_model(main_args)
model.cuda()
model.eval()
images = []
for idx in range(100):
img, t = dataset[idx]
images.append(img)
with torch.no_grad():
tmp = []
for img in tqdm.tqdm(images):
inputs = [img.to('cuda')]
t = measure_time(model, inputs)
tmp.append(t)
res = {'time': fmt_res(np.array(tmp))}
return res
if __name__ == '__main__':
res = benchmark()
print(res)