moe_test.py 935 Bytes
Newer Older
Rick Ho's avatar
Rick Ho committed
1
2
3
from moe import MOELayer
import torch
import time
Rick Ho's avatar
Rick Ho committed
4
import sys
Rick Ho's avatar
Rick Ho committed
5
6
7


def perf():
Rick Ho's avatar
Rick Ho committed
8
9
10
11
    batch_size = int(sys.argv[1])
    in_feat = int(sys.argv[2])
    out_feat = int(sys.argv[3])
    num_expert = int(sys.argv[4])
Rick Ho's avatar
Rick Ho committed
12

Jiezhong Qiu's avatar
Jiezhong Qiu committed
13
14
    inp = torch.rand(batch_size, in_feat).cuda("cuda:1")
    gate = torch.randint(low=0, high=num_expert, size=(batch_size, ), requires_grad=False).int().cuda("cuda:1")
Rick Ho's avatar
Rick Ho committed
15

Jiezhong Qiu's avatar
Jiezhong Qiu committed
16
    moe = MOELayer(num_expert, in_feat, out_feat).cuda("cuda:1")
Rick Ho's avatar
Rick Ho committed
17
18
19
20
21
22

    o = moe(inp, gate)

    n_runs = 16
    tott = 0.
    for i in range(n_runs):
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
23
        gate = torch.randint(low=0, high=num_expert, size=(batch_size, ), requires_grad=False).int().cuda("cuda:1")
Rick Ho's avatar
Rick Ho committed
24
25
26
27
28
        ts = time.time()
        o = moe(inp, gate)
        te = time.time()
        tott += te - ts

Rick Ho's avatar
Rick Ho committed
29
    gflops = 2e-9 * n_runs * in_feat * out_feat * batch_size / tott
Rick Ho's avatar
Rick Ho committed
30
31
32
33
34
    print('Mean time {:.3f} ms, {:.3f} GFLOPs'.format(tott * 1e3 / n_runs, gflops))


if __name__ == '__main__':
    perf()