pytorch

1
import argparse
2
import torch
3

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

9
    def forward(self, x):
10
        y = self.conv(x)
11
        return y
12

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

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