pytorch

1
import argparse
2

3
import torch
4

5

6
class Module(torch.nn.Module):
7
    def __init__(self) -> None:
8
        super().__init__()
9
        self.conv = torch.nn.Conv2d(1, 10, 5, 1)
10

11
    def forward(self, x):
12
        y = self.conv(x)
13
        return y
14

15

16
def run_model(level):
17
    m = Module().eval()
18
    d = torch.rand(1, 1, 112, 112)
19
    with torch.backends.mkldnn.verbose(level):
20
        m(d)
21

22

23
if __name__ == "__main__":
24
    parser = argparse.ArgumentParser()
25
    parser.add_argument("--verbose-level", default=0, type=int)
26
    args = parser.parse_args()
27
    try:
28
        run_model(args.verbose_level)
29
    except Exception as e:
30
        print(e)
31