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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
|
# mypy: allow-untyped-defs
import torch
from torch.nn import (
BatchNorm1d,
BatchNorm2d,
BatchNorm3d,
Conv1d,
Conv2d,
Conv3d,
Linear,
ReLU,
)
from torch.nn.utils.parametrize import type_before_parametrizations
__all__ = [
"ConvReLU1d",
"ConvReLU2d",
"ConvReLU3d",
"LinearReLU",
"ConvBn1d",
"ConvBn2d",
"ConvBnReLU1d",
"ConvBnReLU2d",
"ConvBn3d",
"ConvBnReLU3d",
"BNReLU2d",
"BNReLU3d",
"LinearBn1d",
"LinearLeakyReLU",
"LinearTanh",
"ConvAdd2d",
"ConvAddReLU2d",
]
# Used for identifying intrinsic modules used in quantization
class _FusedModule(torch.nn.Sequential):
pass
class ConvReLU1d(_FusedModule):
r"""This is a sequential container which calls the Conv1d and ReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, relu):
assert (
type_before_parametrizations(conv) == Conv1d
and type_before_parametrizations(relu) == ReLU
), f"Incorrect types for input modules{type_before_parametrizations(conv)}{type_before_parametrizations(relu)}"
super().__init__(conv, relu)
class ConvReLU2d(_FusedModule):
r"""This is a sequential container which calls the Conv2d and ReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, relu):
assert (
type_before_parametrizations(conv) == Conv2d
and type_before_parametrizations(relu) == ReLU
), f"Incorrect types for input modules{type_before_parametrizations(conv)}{type_before_parametrizations(relu)}"
super().__init__(conv, relu)
class ConvReLU3d(_FusedModule):
r"""This is a sequential container which calls the Conv3d and ReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, relu):
assert (
type_before_parametrizations(conv) == Conv3d
and type_before_parametrizations(relu) == ReLU
), f"Incorrect types for input modules{type_before_parametrizations(conv)}{type_before_parametrizations(relu)}"
super().__init__(conv, relu)
class LinearReLU(_FusedModule):
r"""This is a sequential container which calls the Linear and ReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, linear, relu):
assert (
type_before_parametrizations(linear) == Linear
and type_before_parametrizations(relu) == ReLU
), f"Incorrect types for input modules{type_before_parametrizations(linear)}{type_before_parametrizations(relu)}"
super().__init__(linear, relu)
class ConvBn1d(_FusedModule):
r"""This is a sequential container which calls the Conv 1d and Batch Norm 1d modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, bn):
assert (
type_before_parametrizations(conv) == Conv1d
and type_before_parametrizations(bn) == BatchNorm1d
), f"Incorrect types for input modules{type_before_parametrizations(conv)}{type_before_parametrizations(bn)}"
super().__init__(conv, bn)
class ConvBn2d(_FusedModule):
r"""This is a sequential container which calls the Conv 2d and Batch Norm 2d modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, bn):
assert (
type_before_parametrizations(conv) == Conv2d
and type_before_parametrizations(bn) == BatchNorm2d
), f"Incorrect types for input modules{type_before_parametrizations(conv)}{type_before_parametrizations(bn)}"
super().__init__(conv, bn)
class ConvBnReLU1d(_FusedModule):
r"""This is a sequential container which calls the Conv 1d, Batch Norm 1d, and ReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, bn, relu):
assert (
type_before_parametrizations(conv) == Conv1d
and type_before_parametrizations(bn) == BatchNorm1d
and type_before_parametrizations(relu) == ReLU
), f"Incorrect types for input modules{type_before_parametrizations(conv)}{type_before_parametrizations(bn)}{type_before_parametrizations(relu)}" # noqa: B950
super().__init__(conv, bn, relu)
class ConvBnReLU2d(_FusedModule):
r"""This is a sequential container which calls the Conv 2d, Batch Norm 2d, and ReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, bn, relu):
assert (
type_before_parametrizations(conv) == Conv2d
and type_before_parametrizations(bn) == BatchNorm2d
and type_before_parametrizations(relu) == ReLU
), f"Incorrect types for input modules{type_before_parametrizations(conv)}{type_before_parametrizations(bn)}{type_before_parametrizations(relu)}" # noqa: B950
super().__init__(conv, bn, relu)
class ConvBn3d(_FusedModule):
r"""This is a sequential container which calls the Conv 3d and Batch Norm 3d modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, bn):
assert (
type_before_parametrizations(conv) == Conv3d
and type_before_parametrizations(bn) == BatchNorm3d
), f"Incorrect types for input modules{type_before_parametrizations(conv)}{type_before_parametrizations(bn)}"
super().__init__(conv, bn)
class ConvBnReLU3d(_FusedModule):
r"""This is a sequential container which calls the Conv 3d, Batch Norm 3d, and ReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, bn, relu):
assert (
type_before_parametrizations(conv) == Conv3d
and type_before_parametrizations(bn) == BatchNorm3d
and type_before_parametrizations(relu) == ReLU
), f"Incorrect types for input modules{type_before_parametrizations(conv)}{type_before_parametrizations(bn)}{type_before_parametrizations(relu)}" # noqa: B950
super().__init__(conv, bn, relu)
class BNReLU2d(_FusedModule):
r"""This is a sequential container which calls the BatchNorm 2d and ReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, batch_norm, relu):
assert (
type_before_parametrizations(batch_norm) == BatchNorm2d
and type_before_parametrizations(relu) == ReLU
), f"Incorrect types for input modules{type_before_parametrizations(batch_norm)}{type_before_parametrizations(relu)}"
super().__init__(batch_norm, relu)
class BNReLU3d(_FusedModule):
r"""This is a sequential container which calls the BatchNorm 3d and ReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, batch_norm, relu):
assert (
type_before_parametrizations(batch_norm) == BatchNorm3d
and type_before_parametrizations(relu) == ReLU
), f"Incorrect types for input modules{type_before_parametrizations(batch_norm)}{type_before_parametrizations(relu)}"
super().__init__(batch_norm, relu)
class LinearBn1d(_FusedModule):
r"""This is a sequential container which calls the Linear and BatchNorm1d modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, linear, bn):
assert (
type_before_parametrizations(linear) == Linear
and type_before_parametrizations(bn) == BatchNorm1d
), f"Incorrect types for input modules{type_before_parametrizations(linear)}{type_before_parametrizations(bn)}"
super().__init__(linear, bn)
class LinearLeakyReLU(_FusedModule):
r"""This is a sequential container which calls the Linear and LeakyReLU modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, linear, leaky_relu):
assert (
type(linear) == Linear and type(leaky_relu) == torch.nn.LeakyReLU
), f"Incorrect types for input modules{type(linear)}{type(leaky_relu)}"
super().__init__(linear, leaky_relu)
class LinearTanh(_FusedModule):
r"""This is a sequential container which calls the Linear and Tanh modules.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, linear, tanh):
assert (
type(linear) == Linear and type(tanh) == torch.nn.Tanh
), f"Incorrect types for input modules{type(linear)}{type(tanh)}"
super().__init__(linear, tanh)
class ConvAdd2d(_FusedModule):
r"""This is a sequential container which calls the Conv2d modules with extra Add.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, add):
super().__init__(conv)
self.add = add
def forward(self, x1, x2): # type: ignore[override]
return self.add(self[0](x1), x2)
class ConvAddReLU2d(_FusedModule):
r"""This is a sequential container which calls the Conv2d, add, Relu.
During quantization this will be replaced with the corresponding fused module."""
def __init__(self, conv, add, relu):
super().__init__(conv)
self.add = add
self.relu = relu
def forward(self, x1, x2): # type: ignore[override]
return self.relu(self.add(self[0](x1), x2))
|
