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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
|
from collections import namedtuple
from functools import partial
import torchvision.models as cnn
import torch
from .factory import (
dropoutlstm_creator,
imagenet_cnn_creator,
layernorm_pytorch_lstm_creator,
lnlstm_creator,
lstm_creator,
lstm_multilayer_creator,
lstm_premul_bias_creator,
lstm_premul_creator,
lstm_simple_creator,
pytorch_lstm_creator,
varlen_lstm_creator,
varlen_pytorch_lstm_creator,
)
class DisableCuDNN:
def __enter__(self):
self.saved = torch.backends.cudnn.enabled
torch.backends.cudnn.enabled = False
def __exit__(self, *args, **kwargs):
torch.backends.cudnn.enabled = self.saved
class DummyContext:
def __enter__(self):
pass
def __exit__(self, *args, **kwargs):
pass
class AssertNoJIT:
def __enter__(self):
import os
enabled = os.environ.get("PYTORCH_JIT", 1)
assert not enabled
def __exit__(self, *args, **kwargs):
pass
RNNRunner = namedtuple(
"RNNRunner",
[
"name",
"creator",
"context",
],
)
def get_nn_runners(*names):
return [nn_runners[name] for name in names]
nn_runners = {
"cudnn": RNNRunner("cudnn", pytorch_lstm_creator, DummyContext),
"cudnn_dropout": RNNRunner(
"cudnn_dropout", partial(pytorch_lstm_creator, dropout=0.4), DummyContext
),
"cudnn_layernorm": RNNRunner(
"cudnn_layernorm", layernorm_pytorch_lstm_creator, DummyContext
),
"vl_cudnn": RNNRunner("vl_cudnn", varlen_pytorch_lstm_creator, DummyContext),
"vl_jit": RNNRunner(
"vl_jit", partial(varlen_lstm_creator, script=True), DummyContext
),
"vl_py": RNNRunner("vl_py", varlen_lstm_creator, DummyContext),
"aten": RNNRunner("aten", pytorch_lstm_creator, DisableCuDNN),
"jit": RNNRunner("jit", lstm_creator, DummyContext),
"jit_premul": RNNRunner("jit_premul", lstm_premul_creator, DummyContext),
"jit_premul_bias": RNNRunner(
"jit_premul_bias", lstm_premul_bias_creator, DummyContext
),
"jit_simple": RNNRunner("jit_simple", lstm_simple_creator, DummyContext),
"jit_multilayer": RNNRunner(
"jit_multilayer", lstm_multilayer_creator, DummyContext
),
"jit_layernorm": RNNRunner("jit_layernorm", lnlstm_creator, DummyContext),
"jit_layernorm_decom": RNNRunner(
"jit_layernorm_decom",
partial(lnlstm_creator, decompose_layernorm=True),
DummyContext,
),
"jit_dropout": RNNRunner("jit_dropout", dropoutlstm_creator, DummyContext),
"py": RNNRunner("py", partial(lstm_creator, script=False), DummyContext),
"resnet18": RNNRunner(
"resnet18", imagenet_cnn_creator(cnn.resnet18, jit=False), DummyContext
),
"resnet18_jit": RNNRunner(
"resnet18_jit", imagenet_cnn_creator(cnn.resnet18), DummyContext
),
"resnet50": RNNRunner(
"resnet50", imagenet_cnn_creator(cnn.resnet50, jit=False), DummyContext
),
"resnet50_jit": RNNRunner(
"resnet50_jit", imagenet_cnn_creator(cnn.resnet50), DummyContext
),
}
|
