intel-extension-for-pytorch

1
import torch
2
import intel_extension_for_pytorch as ipex
3
from common_utils import int8_calibration
4

5
ipex.core.enable_auto_dnnl()
6

7
ic = 1024
8
oc = 1024
9
bs = 16
10

11
LL = torch.nn.Linear(ic, oc).to(ipex.DEVICE)
12

13

14
def get_input():
15
    return torch.rand(bs, ic).to(ipex.DEVICE)
16

17

18
def run_linear(auto_mix_conf=None):
19
    for i in range(3):
20
        if auto_mix_conf is not None:
21
            with ipex.AutoMixPrecision(auto_mix_conf):
22
                LL(get_input())
23
        else:
24
            LL(get_input())
25

26

27
if __name__ == "__main__":
28
    print(f"fp32, {'*' * 50}")
29
    run_linear()
30

31
    print(f"auto-mix for bf16, {'*' * 50}")
32
    bf16_conf = ipex.AmpConf(torch.bfloat16)
33
    run_linear(bf16_conf)
34

35
    print(f"back to fp32, {'*' * 50}")
36
    ipex.core.reorder_to_float32(LL.weight)
37
    ipex.core.reorder_to_float32(LL.bias)
38
    run_linear()
39

40
    print(f"auto-mix for int8, {'*' * 50}")
41
    int8_calibration(LL, [get_input() for i in range(3)], "./int8.config")
42
    int8_conf = ipex.AmpConf(torch.int8, "./int8.config")
43
    run_linear(int8_conf)
44

45
    print(f"back to fp32, {'*' * 50}")
46
    ipex.core.reorder_to_float32(LL.weight)
47
    ipex.core.reorder_to_float32(LL.bias)
48
    run_linear()
49