forked from seba-1511/lstms.pth
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_speed.py
70 lines (61 loc) · 2.2 KB
/
test_speed.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
#!/usr/bin/env python
import torch as th
import torch.nn as nn
from torch.autograd import Variable as V
from time import time
from lstms import (SlowLSTM,
LSTM,
GalLSTM,
MoonLSTM,
SemeniutaLSTM,
LayerNormLSTM,
LayerNormGalLSTM,
LayerNormMoonLSTM,
LayerNormSemeniutaLSTM,
)
if __name__ == '__main__':
N_ITER = 1000
SIZES = [128, 256, 512, 1024, 2048]
DROPOUT = 0.5
lstms = [
(SlowLSTM, 'SlowLSTM'),
(LSTM, 'LSTM'),
(GalLSTM, 'GalLSTM'),
(MoonLSTM, 'MoonLSTM'),
(SemeniutaLSTM, 'SemeniutaLSTM'),
(LayerNormLSTM, 'LayerNormLSTM'),
(LayerNormGalLSTM, 'LayerNormGalLSTM'),
(LayerNormMoonLSTM, 'LayerNormMoonLSTM'),
(LayerNormSemeniutaLSTM, 'LayerNormSemeniutaLSTM'),
]
for lstm, name in lstms:
ref_res = []
cus_res = []
for size in SIZES:
x = V(th.rand(1, 1, size))
hiddens = (V(th.rand(1, 1, size)), V(th.rand(1, 1, size)))
th.manual_seed(1234)
ref = nn.LSTM(size, size, dropout=DROPOUT)
th.manual_seed(1234)
cus = lstm(size, size, dropout=DROPOUT)
out, h = x, hiddens
ref_start = time()
for i in range(N_ITER):
out, h = ref(out, h)
ref_time = time() - ref_start
ref_res.append(ref_time)
out, h = x, hiddens
cus_start = time()
for i in range(N_ITER):
out, h = cus(out, h)
cus_time = time() - cus_start
cus_res.append(cus_time)
print('## ', name, ' Benchmark ')
print(' ')
print('Inference timings on a single sequence of length', N_ITER, ' with `dropout = ', DROPOUT, '`.')
print(' ')
print('size | nn.LSTM | ', name, ' | Speedup')
print('-------|-----------|-' + '-' * len(name) + '---|--------')
for size, ref, cus in zip(SIZES, ref_res, cus_res):
print(size, (' | %.3f | %.3f' + ' ' * (len(name) - 4) + ' | %.3f') % (ref, cus, ref / cus))
print(' ')